python3Packages.dm-haiku: run tests in parallel

2022-03-29 18:13:47 +00:00 · 2022-03-29 18:13:47 +00:00 · 740a83a9bd
parent 3c25104d0e
commit 740a83a9bd
1 changed files with 3 additions and 0 deletions
--- a/pkgs/development/python-modules/dm-haiku/default.nix
+++ b/pkgs/development/python-modules/dm-haiku/default.nix
@ -7,6 +7,7 @@
 , jaxlib
 , jmp
 , lib
+, pytest-xdist
 , pytestCheckHook
 , tabulate
 , tensorflow
@ -33,9 +34,11 @@ buildPythonPackage rec {
    cloudpickle
    dm-tree
    jaxlib
+    pytest-xdist
    pytestCheckHook
    tensorflow
  ];
+  pytestFlagsArray = [ "-n $NIX_BUILD_CORES" ];

  pythonImportsCheck = [
    "haiku"