From d57404ea3dcb3ab98913414a6f6fba854a8509b2 Mon Sep 17 00:00:00 2001
From: Martin Weinelt <hexa@darmstadt.ccc.de>
Date: Sun, 10 Apr 2022 17:31:15 +0200
Subject: [PATCH] python3Packages.jax: test with limited parallelism

The tests are prone to getting stuck with high parallelism.
---
 pkgs/development/python-modules/jax/default.nix | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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