From 4570907282908a9f423807750994212abd803412 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Tue, 29 Mar 2022 18:14:04 +0000
Subject: [PATCH] python3Packages.optax: run tests in parallel

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

diff --git a/pkgs/development/python-modules/optax/default.nix b/pkgs/development/python-modules/optax/default.nix
index 6a3b6a9d3e67..ae8cfdbbc282 100644
--- a/pkgs/development/python-modules/optax/default.nix
+++ b/pkgs/development/python-modules/optax/default.nix
@@ -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"