Revert "python3Packages.jax: remove pytest-xdist"

This reverts commit bddafba6a5.

Running without pytest-xdist leads to reliable crashes on both
aarch64-linux (Ampere Q80-30) and x86_64-linux (Zen 3 Epyc) on hydra,
which prevents the package from being cached.

We revert this change to prevent the churn on downstream maintainers
and consumers, that would otherwise have to build this expensive
package locally, which does not work reliably.

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/"
   ];