diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index ec9b429b8931..e82e53219bc3 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -102,19 +102,15 @@ buildPythonPackage rec {
     "flax/nnx/examples/*"
     # See https://github.com/google/flax/issues/3232.
     "tests/jax_utils_test.py"
-    # Too old version of tensorflow:
-    # ModuleNotFoundError: No module named 'keras.api._v2'
-    "tests/tensorboard_test.py"
   ];
 
   disabledTests =
     [
-      # ValueError: Checkpoint path should be absolute
-      "test_overwrite_checkpoints0"
-      # Fixed in more recent versions of jax: https://github.com/google/flax/issues/4211
-      # TODO: Re-enable when jax>0.4.28 will be available in nixpkgs
-      "test_vmap_and_cond_passthrough" # ValueError: vmap has mapped output but out_axes is None
-      "test_vmap_and_cond_passthrough_error" # AssertionError: "at vmap.*'broadcast'.*got axis spec ...
+      # Failing with AssertionError since the jax update to 0.5.0
+      "test_basic_demo_single"
+      "test_batch_norm_multi_init"
+      "test_multimetric"
+      "test_split_merge"
     ]
     ++ lib.optionals stdenv.hostPlatform.isDarwin [
       # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!