From 6fa30f68552aa526a13382f2f0120736660219d2 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Tue, 29 Mar 2022 18:13:57 +0000
Subject: [PATCH] python3Packages.flax: run tests in parallel

---
 pkgs/development/python-modules/flax/default.nix | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index 3ac8166bbbd..1f1622dc1ae 100644
--- 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.