pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 flax,
5 tomlq,
6 python,
7
8 # build-system
9 nanobind,
10 ninja,
11 scikit-build-core,
12
13 # nativeBuildInputs
14 cmake,
15 pkg-config,
16}:
17
18buildPythonPackage rec {
19 pname = "flaxlib";
20 version = "0.0.1";
21 pyproject = true;
22
23 inherit (flax) src;
24
25 sourceRoot = "${src.name}/flaxlib_src";
26
27 postPatch = ''
28 expected_version="$version"
29 actual_version=$(${lib.getExe tomlq} --raw --file pyproject.toml "project.version")
30
31 if [ "$actual_version" != "$expected_version" ]; then
32 echo -e "\n\tERROR:"
33 echo -e "\tThe version of the flaxlib python package ($expected_version) does not match the one in its pyproject.toml file ($actual_version)"
34 echo -e "\tPlease update the version attribute of the nix python3Packages.flaxlib package."
35 exit 1
36 fi
37 '';
38
39 dontUseCmakeConfigure = true;
40
41 build-system = [
42 nanobind
43 ninja
44 scikit-build-core
45 ];
46 nativeBuildInputs = [
47 cmake
48 pkg-config
49 ];
50
51 env.CMAKE_PREFIX_PATH = "${nanobind}/${python.sitePackages}/nanobind";
52
53 pythonImportsCheck = [ "flaxlib" ];
54
55 # This package does not have tests (yet ?)
56 doCheck = false;
57
58 passthru = {
59 inherit (flax) updateScript;
60 };
61
62 meta = {
63 description = "Rust library used internally by flax";
64 homepage = "https://github.com/google/flax/tree/main/flaxlib";
65 license = lib.licenses.asl20;
66 maintainers = with lib.maintainers; [ GaetanLepage ];
67 };
68}