Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Build jax/jaxlib from source against MKL version 2021.1.1 and run the complete jax test suite.

In nix, this can be accomplished relatively easily:

Checkout nixpkgs at NixOS/nixpkgs@c38d365.

Add

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; };
}

to ~/.config/nixpkgs/overlays/cuda_and_mkl.nix.

nix-build -A python3Packages.jax

If applicable, include full error messages/tracebacks.

=================================== FAILURES ===================================
___________________ CustomRootTest.test_custom_root_with_aux ___________________
[gw8] linux -- Python 3.9.10 /nix/store/afi0ysqw20yiiw2gr2d28dx40bc4ddf8-python3-3.9.10/bin/python3.9

self = <custom_root_test.CustomRootTest testMethod=test_custom_root_with_aux>

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)
> 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.
> 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;

What versions of BLAS/LAPACK does JAX support? Is using MKL with JAX an officially supported combination?

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

import jax.numpy as jnp
print(jnp.arange(10))
# [0 1 2 3 4 5 6 7 8 9]

If applicable, include full error messages/tracebacks.

I am trying to build cuda-enabled jaxlib on conda-forge. I am running into a problem where the --config=cuda 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 --config=cuda works fine in our tensorflow builds.

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

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.

Corresponding PR in conda-forge: conda-forge/jaxlib-feedstock#97

Tagging @hawkinsp who's been involved in issues related to the builds before. Thanks!