python311Packages.flax: refactor

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

@@ -1,24 +1,26 @@
 { lib
 , buildPythonPackage
+, cloudpickle
+, einops
 , fetchFromGitHub
-, jaxlib
-, pythonRelaxDepsHook
-, setuptools-scm
 , jax
+, jaxlib
+, keras
+, matplotlib
 , msgpack
 , numpy
 , optax
-, pyyaml
-, rich
-, tensorstore
-, typing-extensions
-, matplotlib
-, cloudpickle
-, einops
-, keras
+, orbax-checkpoint
 , pytest-xdist
 , pytestCheckHook
+, pythonOlder
+, pythonRelaxDepsHook
+, pyyaml
+, rich
+, setuptools-scm
 , tensorflow
+, tensorstore
+, typing-extensions
 }:
 
 buildPythonPackage rec {
@@ -26,6 +28,8 @@ buildPythonPackage rec {
   version = "0.7.5";
   pyproject = true;
 
+  disabled = pythonOlder "3.8";
+
   src = fetchFromGitHub {
     owner = "google";
     repo = "flax";
@@ -44,6 +48,7 @@ buildPythonPackage rec {
     msgpack
     numpy
     optax
+    orbax-checkpoint
     pyyaml
     rich
     tensorstore
@@ -75,7 +80,6 @@ buildPythonPackage rec {
   disabledTestPaths = [
     # Docs test, needs extra deps + we're not interested in it.
     "docs/_ext/codediff_test.py"
-
     # The tests in `examples` are not designed to be executed from a single test
     # session and thus either have the modules that conflict with each other or
     # wrong import paths, depending on how they're invoked. Many tests also have
@@ -83,12 +87,15 @@ buildPythonPackage rec {
     # `tensorflow_datasets`, `vocabulary`) so the benefits of trying to run them
     # would be limited anyway.
     "examples/*"
-
     # See https://github.com/google/flax/issues/3232.
     "tests/jax_utils_test.py"
+    # Requires tree
+    "tests/tensorboard_test.py"
+  ];
 
-    # Requires orbax which is not packaged as of 2023-07-27.
-    "tests/checkpoints_test.py"
+  disabledTests = [
+    # ValueError: Checkpoint path should be absolute
+    "test_overwrite_checkpoints0"
   ];
 
   meta = with lib; {