1 files changed, 15 insertions, 0 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/jax/default.nix b/nixpkgs/pkgs/development/python-modules/jax/default.nix
index c91b6e485226..d5e536540195 100644
--- a/nixpkgs/pkgs/development/python-modules/jax/default.nix
+++ b/nixpkgs/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,16 @@ buildPythonPackage rec {
     "tests/"
   ];
 
+  # See
+  #  * https://github.com/google/jax/issues/9705
+  #  * https://discourse.nixos.org/t/getting-different-results-for-the-same-build-on-two-equally-configured-machines/17921
+  #  * https://github.com/NixOS/nixpkgs/issues/161960
+  disabledTests = lib.optionals usingMKL [
+    "test_custom_linear_solve_cholesky"
+    "test_custom_root_with_aux"
+    "testEigvalsGrad_shape"
+  ];
+
   pythonImportsCheck = [
     "jax"
   ];