python3Packages.flax: run tests in parallel

2 years ago · 6fa30f6855
parent 740a83a9bd
commit 6fa30f6855
1 changed files with 3 additions and 0 deletions
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@ -7,6 +7,7 @@
 , msgpack
 , numpy
 , optax
+, pytest-xdist
 , pytestCheckHook
 , tensorflow
 }:
@ -37,9 +38,11 @@ buildPythonPackage rec {

  checkInputs = [
    keras
+    pytest-xdist
    pytestCheckHook
    tensorflow
  ];
+  pytestFlagsArray = [ "-n $NIX_BUILD_CORES" ];

  disabledTestPaths = [
    # Docs test, needs extra deps + we're not interested in it.