diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index 332cac3133cc..0d1e8c17b13e 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -54,12 +54,15 @@ buildPythonPackage rec { pytest-xdist ]; + # high parallelism will result in the tests getting stuck + dontUsePytestXdist = true; + # NOTE: Don't run the tests in the expiremental directory as they require flax # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2. # Not a big deal, this is how the JAX docs suggest running the test suite # anyhow. pytestFlagsArray = [ - "-n auto" + "--numprocesses=4" "-W ignore::DeprecationWarning" "tests/" ];