pkgs/development/python-modules/flax/default.nix

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchFromGitHub,
   5
   6   # build-system
   7   jaxlib,
   8   setuptools-scm,
   9
  10   # dependencies
  11   jax,
  12   msgpack,
  13   numpy,
  14   optax,
  15   orbax-checkpoint,
  16   pyyaml,
  17   rich,
  18   tensorstore,
  19   typing-extensions,
  20
  21   # checks
  22   cloudpickle,
  23   einops,
  24   flaxlib,
  25   keras,
  26   pytestCheckHook,
  27   pytest-xdist,
  28   sphinx,
  29   tensorflow,
  30   treescope,
  31
  32   # optional-dependencies
  33   matplotlib,
  34
  35   writeScript,
  36   tomlq,
  37 }:
  38
  39 buildPythonPackage rec {
  40   pname = "flax";
  41   version = "0.10.1";
  42   pyproject = true;
  43
  44   src = fetchFromGitHub {
  45     owner = "google";
  46     repo = "flax";
  47     rev = "refs/tags/v${version}";
  48     hash = "sha256-+URbQGnmqmSNgucEyWvI5DMnzXjpmJzLA+Pho2lX+S4=";
  49   };
  50
  51   build-system = [
  52     jaxlib
  53     setuptools-scm
  54   ];
  55
  56   dependencies = [
  57     jax
  58     msgpack
  59     numpy
  60     optax
  61     orbax-checkpoint
  62     pyyaml
  63     rich
  64     tensorstore
  65     typing-extensions
  66   ];
  67
  68   optional-dependencies = {
  69     all = [ matplotlib ];
  70   };
  71
  72   pythonImportsCheck = [ "flax" ];
  73
  74   nativeCheckInputs = [
  75     cloudpickle
  76     einops
  77     flaxlib
  78     keras
  79     pytestCheckHook
  80     pytest-xdist
  81     sphinx
  82     tensorflow
  83     treescope
  84   ];
  85
  86   pytestFlagsArray = [
  87     "-W ignore::FutureWarning"
  88     "-W ignore::DeprecationWarning"
  89   ];
  90
  91   disabledTestPaths = [
  92     # Docs test, needs extra deps + we're not interested in it.
  93     "docs/_ext/codediff_test.py"
  94     # The tests in `examples` are not designed to be executed from a single test
  95     # session and thus either have the modules that conflict with each other or
  96     # wrong import paths, depending on how they're invoked. Many tests also have
  97     # dependencies that are not packaged in `nixpkgs` (`clu`, `jgraph`,
  98     # `tensorflow_datasets`, `vocabulary`) so the benefits of trying to run them
  99     # would be limited anyway.
 100     "examples/*"
 101     "flax/nnx/examples/*"
 102     # See https://github.com/google/flax/issues/3232.
 103     "tests/jax_utils_test.py"
 104     # Too old version of tensorflow:
 105     # ModuleNotFoundError: No module named 'keras.api._v2'
 106     "tests/tensorboard_test.py"
 107   ];
 108
 109   disabledTests = [
 110     # ValueError: Checkpoint path should be absolute
 111     "test_overwrite_checkpoints0"
 112     # Fixed in more recent versions of jax: https://github.com/google/flax/issues/4211
 113     # TODO: Re-enable when jax>0.4.28 will be available in nixpkgs
 114     "test_vmap_and_cond_passthrough" # ValueError: vmap has mapped output but out_axes is None
 115     "test_vmap_and_cond_passthrough_error" # AssertionError: "at vmap.*'broadcast'.*got axis spec ...
 116   ];
 117
 118   passthru = {
 119     updateScript = writeScript "update.sh" ''
 120       nix-update flax # does not --build by default
 121       nix-build . -A flax.src # src is essentially a passthru
 122       nix-update flaxlib --version="$(${lib.getExe tomlq} <result/Cargo.toml .something.version)" --commit
 123     '';
 124   };
 125
 126   meta = {
 127     description = "Neural network library for JAX";
 128     homepage = "https://github.com/google/flax";
 129     changelog = "https://github.com/google/flax/releases/tag/v${version}";
 130     license = lib.licenses.asl20;
 131     maintainers = with lib.maintainers; [ ndl ];
 132   };
 133 }