python3Packages.optax: run tests in parallel

2022-03-30 21:34:09 +02:00
parent 6fa30f6855
commit 4570907282
1 changed files with 3 additions and 0 deletions
@@ -6,6 +6,7 @@
 , jaxlib
 , lib
 , numpy
+, pytest-xdist
 , pytestCheckHook
 }:

@@ -32,8 +33,10 @@ buildPythonPackage rec {

  checkInputs = [
    dm-haiku
+    pytest-xdist
    pytestCheckHook
  ];
+  pytestFlagsArray = [ "-n $NIX_BUILD_CORES" ];

  pythonImportsCheck = [
    "optax"