diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index dbdf655e4007..25ff005569a5 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -12,6 +12,7 @@ , numpy , opt-einsum , pytestCheckHook +, pytest-xdist , pythonOlder , scipy , stdenv @@ -61,13 +62,18 @@ buildPythonPackage rec { jaxlib' matplotlib pytestCheckHook + 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 = [ + "--numprocesses=4" "-W ignore::DeprecationWarning" "tests/" ];