<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">Build jax/jaxlib from source against MKL version 2021.1.1 and run the complete jax test suite.</p>
<p dir="auto">In nix, this can be accomplished relatively easily:</p>
<ol dir="auto">
<li>Checkout nixpkgs at <a class="commit-link" data-hovercard-type="commit" data-hovercard-url="https://github.com/NixOS/nixpkgs/commit/c38d3659e3fa2e63055197b8b73bba0cd55a40dc/hovercard" href="https://github.com/NixOS/nixpkgs/commit/c38d3659e3fa2e63055197b8b73bba0cd55a40dc">NixOS/nixpkgs@<tt>c38d365</tt></a>.</li>
<li>Add</li>
</ol>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="final: prev: {
  cudatoolkit = prev.cudatoolkit_11_5;
  cudnn = prev.cudnn_8_3_cudatoolkit_11_5;
  blas = prev.blas.override { blasProvider = final.mkl; };
  lapack = prev.lapack.override { lapackProvider = final.mkl; };
}"><pre class="notranslate"><code class="notranslate">final: prev: {
  cudatoolkit = prev.cudatoolkit_11_5;
  cudnn = prev.cudnn_8_3_cudatoolkit_11_5;
  blas = prev.blas.override { blasProvider = final.mkl; };
  lapack = prev.lapack.override { lapackProvider = final.mkl; };
}
</code></pre></div>
<p dir="auto">to <code class="notranslate">~/.config/nixpkgs/overlays/cuda_and_mkl.nix</code>.</p>
<ol start="3" dir="auto">
<li><code class="notranslate">nix-build -A python3Packages.jax</code></li>
</ol>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="=================================== FAILURES ===================================
___________________ CustomRootTest.test_custom_root_with_aux ___________________
[gw8] linux -- Python 3.9.10 /nix/store/afi0ysqw20yiiw2gr2d28dx40bc4ddf8-python3-3.9.10/bin/python3.9

self = &lt;custom_root_test.CustomRootTest testMethod=test_custom_root_with_aux&gt;

    def test_custom_root_with_aux(self):
      def root_aux(a, b):
        f = lambda x: high_precision_dot(a, x) - b
        factors = jsp.linalg.cho_factor(a)
        cho_solve = lambda f, b: (jsp.linalg.cho_solve(factors, b), orig_aux)
    
        def pos_def_solve(g, b):
          # prune aux to allow use as tangent_solve
          cho_solve_noaux = lambda f, b: cho_solve(f, b)[0]
          return lax.custom_linear_solve(g, b, cho_solve_noaux, symmetric=True)
    
        return lax.custom_root(f, b, cho_solve, pos_def_solve, has_aux=True)
    
      orig_aux = {&quot;converged&quot;: np.array(1.), &quot;nfev&quot;: np.array(12345.), &quot;grad&quot;: np.array([1.0, 2.0, 3.0])}
    
      rng = self.rng()
      a = rng.randn(2, 2)
      b = rng.randn(2)
    
      actual, actual_aux = root_aux(high_precision_dot(a, a.T), b)
      actual_jit, actual_jit_aux = jax.jit(root_aux)(high_precision_dot(a, a.T), b)
      expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    
      self.assertAllClose(expected, actual)
      self.assertAllClose(expected, actual_jit)
      jtu.check_eq(actual_jit_aux, orig_aux)
    
      # grad check with aux
      jtu.check_grads(lambda x, y: root_aux(high_precision_dot(x, x.T), y),
                      (a, b), order=2, rtol={jnp.float32: 1e-2})
    
      # test vmap and jvp combined by jacfwd
      fwd = jax.jacfwd(lambda x, y: root_aux(high_precision_dot(x, x.T), y), argnums=(0, 1))
      expected_fwd = jax.jacfwd(lambda x, y: jnp.linalg.solve(high_precision_dot(x, x.T), y), argnums=(0, 1))
    
      fwd_val, fwd_aux = fwd(a, b)
      expected_fwd_val = expected_fwd(a, b)
&gt;     self.assertAllClose(fwd_val, expected_fwd_val, rtol={np.float32: 1E-6, np.float64: 1E-12})

tests/custom_root_test.py:227: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
jax/_src/test_util.py:990: in assertAllClose
    self.assertAllClose(x_elt, y_elt, check_dtypes=check_dtypes, atol=atol,
jax/_src/test_util.py:999: in assertAllClose
    self.assertArraysAllClose(x, y, check_dtypes=False, atol=atol, rtol=rtol,
jax/_src/test_util.py:964: in assertArraysAllClose
    _assert_numpy_allclose(x, y, atol=atol, rtol=rtol, err_msg=err_msg)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

a = array([[[-1.3422977 ,  0.55962217],
        [ 1.0310359 , -0.8444528 ]],

       [[ 1.4493693 , -0.3698212 ],
        [-1.4489349 ,  1.1834809 ]]], dtype=float32)
b = array([[[-1.3422956 ,  0.5596222 ],
        [ 1.0310338 , -0.8444524 ]],

       [[ 1.4493668 , -0.36982143],
        [-1.4489323 ,  1.1834806 ]]], dtype=float32)
atol = 1e-06, rtol = 1e-06, err_msg = ''

    def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
      if a.dtype == b.dtype == _dtypes.float0:
        np.testing.assert_array_equal(a, b, err_msg=err_msg)
        return
      a = a.astype(np.float32) if a.dtype == _dtypes.bfloat16 else a
      b = b.astype(np.float32) if b.dtype == _dtypes.bfloat16 else b
      kw = {}
      if atol: kw[&quot;atol&quot;] = atol
      if rtol: kw[&quot;rtol&quot;] = rtol
      with np.errstate(invalid='ignore'):
        # TODO(phawkins): surprisingly, assert_allclose sometimes reports invalid
        # value errors. It should not do that.
&gt;       np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
E       AssertionError: 
E       Not equal to tolerance rtol=1e-06, atol=1e-06
E       
E       Mismatched elements: 3 / 8 (37.5%)
E       Max absolute difference: 2.6226044e-06
E       Max relative difference: 2.0811804e-06
E        x: array([[[-1.342298,  0.559622],
E               [ 1.031036, -0.844453]],
E       ...
E        y: array([[[-1.342296,  0.559622],
E               [ 1.031034, -0.844452]],
E       ...

jax/_src/test_util.py:185: AssertionError
=========================== short test summary info ============================
FAILED tests/custom_root_test.py::CustomRootTest::test_custom_root_with_aux
========== 1 failed, 13955 passed, 1409 skipped in 216.44s (0:03:36) ===========
error: builder for '/nix/store/7z1ydr7xc19108r3nlp5na4m45izascl-python3.9-jax-0.3.1.drv' failed with exit code 1;"><pre class="notranslate"><code class="notranslate">=================================== FAILURES ===================================
___________________ CustomRootTest.test_custom_root_with_aux ___________________
[gw8] linux -- Python 3.9.10 /nix/store/afi0ysqw20yiiw2gr2d28dx40bc4ddf8-python3-3.9.10/bin/python3.9

self = &lt;custom_root_test.CustomRootTest testMethod=test_custom_root_with_aux&gt;

    def test_custom_root_with_aux(self):
      def root_aux(a, b):
        f = lambda x: high_precision_dot(a, x) - b
        factors = jsp.linalg.cho_factor(a)
        cho_solve = lambda f, b: (jsp.linalg.cho_solve(factors, b), orig_aux)
    
        def pos_def_solve(g, b):
          # prune aux to allow use as tangent_solve
          cho_solve_noaux = lambda f, b: cho_solve(f, b)[0]
          return lax.custom_linear_solve(g, b, cho_solve_noaux, symmetric=True)
    
        return lax.custom_root(f, b, cho_solve, pos_def_solve, has_aux=True)
    
      orig_aux = {"converged": np.array(1.), "nfev": np.array(12345.), "grad": np.array([1.0, 2.0, 3.0])}
    
      rng = self.rng()
      a = rng.randn(2, 2)
      b = rng.randn(2)
    
      actual, actual_aux = root_aux(high_precision_dot(a, a.T), b)
      actual_jit, actual_jit_aux = jax.jit(root_aux)(high_precision_dot(a, a.T), b)
      expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    
      self.assertAllClose(expected, actual)
      self.assertAllClose(expected, actual_jit)
      jtu.check_eq(actual_jit_aux, orig_aux)
    
      # grad check with aux
      jtu.check_grads(lambda x, y: root_aux(high_precision_dot(x, x.T), y),
                      (a, b), order=2, rtol={jnp.float32: 1e-2})
    
      # test vmap and jvp combined by jacfwd
      fwd = jax.jacfwd(lambda x, y: root_aux(high_precision_dot(x, x.T), y), argnums=(0, 1))
      expected_fwd = jax.jacfwd(lambda x, y: jnp.linalg.solve(high_precision_dot(x, x.T), y), argnums=(0, 1))
    
      fwd_val, fwd_aux = fwd(a, b)
      expected_fwd_val = expected_fwd(a, b)
&gt;     self.assertAllClose(fwd_val, expected_fwd_val, rtol={np.float32: 1E-6, np.float64: 1E-12})

tests/custom_root_test.py:227: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
jax/_src/test_util.py:990: in assertAllClose
    self.assertAllClose(x_elt, y_elt, check_dtypes=check_dtypes, atol=atol,
jax/_src/test_util.py:999: in assertAllClose
    self.assertArraysAllClose(x, y, check_dtypes=False, atol=atol, rtol=rtol,
jax/_src/test_util.py:964: in assertArraysAllClose
    _assert_numpy_allclose(x, y, atol=atol, rtol=rtol, err_msg=err_msg)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

a = array([[[-1.3422977 ,  0.55962217],
        [ 1.0310359 , -0.8444528 ]],

       [[ 1.4493693 , -0.3698212 ],
        [-1.4489349 ,  1.1834809 ]]], dtype=float32)
b = array([[[-1.3422956 ,  0.5596222 ],
        [ 1.0310338 , -0.8444524 ]],

       [[ 1.4493668 , -0.36982143],
        [-1.4489323 ,  1.1834806 ]]], dtype=float32)
atol = 1e-06, rtol = 1e-06, err_msg = ''

    def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
      if a.dtype == b.dtype == _dtypes.float0:
        np.testing.assert_array_equal(a, b, err_msg=err_msg)
        return
      a = a.astype(np.float32) if a.dtype == _dtypes.bfloat16 else a
      b = b.astype(np.float32) if b.dtype == _dtypes.bfloat16 else b
      kw = {}
      if atol: kw["atol"] = atol
      if rtol: kw["rtol"] = rtol
      with np.errstate(invalid='ignore'):
        # TODO(phawkins): surprisingly, assert_allclose sometimes reports invalid
        # value errors. It should not do that.
&gt;       np.testing.assert_allclose(a, b, **kw, err_msg=err_msg)
E       AssertionError: 
E       Not equal to tolerance rtol=1e-06, atol=1e-06
E       
E       Mismatched elements: 3 / 8 (37.5%)
E       Max absolute difference: 2.6226044e-06
E       Max relative difference: 2.0811804e-06
E        x: array([[[-1.342298,  0.559622],
E               [ 1.031036, -0.844453]],
E       ...
E        y: array([[[-1.342296,  0.559622],
E               [ 1.031034, -0.844452]],
E       ...

jax/_src/test_util.py:185: AssertionError
=========================== short test summary info ============================
FAILED tests/custom_root_test.py::CustomRootTest::test_custom_root_with_aux
========== 1 failed, 13955 passed, 1409 skipped in 216.44s (0:03:36) ===========
error: builder for '/nix/store/7z1ydr7xc19108r3nlp5na4m45izascl-python3.9-jax-0.3.1.drv' failed with exit code 1;
</code></pre></div>
<p dir="auto">What versions of BLAS/LAPACK does JAX support? Is using MKL with JAX an officially supported combination?</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
print(jnp.arange(10))
# [0 1 2 3 4 5 6 7 8 9]"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-en">print</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">10</span>))
<span class="pl-c"># [0 1 2 3 4 5 6 7 8 9]</span></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<hr>
<p dir="auto">I am trying to build cuda-enabled jaxlib on conda-forge. I am running into a problem where the <code class="notranslate">--config=cuda</code> flag replaces the custom toolchain. I understand this is outside the scope here, but getting cuda-enabled jax on conda-forge seems to be a very valuable contribution, and I hope you can help us achieve that. We do build cuda-enabled tensorflow, and we follow our method almost exactly. The option <code class="notranslate">--config=cuda</code> works fine in our tensorflow builds.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="WARNING: option '--config=cuda' (source command line options) was expanded and now overrides the explicit option --crosstool_top=//bazel_toolchain:toolchain with --crosstool_top=@local_config_cuda//crosstool:toolchain"><pre class="notranslate"><code class="notranslate">WARNING: option '--config=cuda' (source command line options) was expanded and now overrides the explicit option --crosstool_top=//bazel_toolchain:toolchain with --crosstool_top=@local_config_cuda//crosstool:toolchain
</code></pre></div>
<p dir="auto">Once the above happens, the local_config_cuda parameters essentially cannot see any of the system libraries and everything  breaks. I haven't looked too deeply into this yet, so bringing it here first in case you have a quick answer/shortcut.</p>
<p dir="auto">Corresponding PR in conda-forge: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="1234482537" data-permission-text="Title is private" data-url="https://github.com/conda-forge/jaxlib-feedstock/issues/97" data-hovercard-type="pull_request" data-hovercard-url="/conda-forge/jaxlib-feedstock/pull/97/hovercard" href="https://github.com/conda-forge/jaxlib-feedstock/pull/97">conda-forge/jaxlib-feedstock#97</a></p>
<p dir="auto">Tagging <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/hawkinsp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/hawkinsp">@hawkinsp</a> who's been involved in issues related to the builds before. Thanks!</p>