diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index c91b6e485226..a616a83255b6 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -1,8 +1,10 @@ { lib , absl-py +, blas , buildPythonPackage , fetchFromGitHub , jaxlib +, lapack , numpy , opt-einsum , pytestCheckHook @@ -12,6 +14,9 @@ , typing-extensions }: +let + usingMKL = blas.implementation == "mkl" || lapack.implementation == "mkl"; +in buildPythonPackage rec { pname = "jax"; version = "0.3.1"; @@ -59,6 +64,9 @@ buildPythonPackage rec { "tests/" ]; + # See https://github.com/google/jax/issues/9705. + disabledTests = lib.optionals usingMKL [ "test_custom_root_with_aux" ]; + pythonImportsCheck = [ "jax" ];