Repository: google/jax

Issue:
jax.numpy.linalg ops missing GPU implementation
I often hit messages like 
```"Singular value decomposition is only implemented on the CPU backend"```
(https://github.com/google/jax/blob/27746b8c73f9ca9928da5da40b7382ae648a5f8d/jax/lax_linalg.py)
for many `jax.numpy.linalg` ops. Examples I've hit so far are when calling:
`jax.numpy.linalg.svd`;
`jax.numpy.linalg.eigh`.

It would be nice to be able to run on GPU.

Thanks!


Discussion:
Thanks for raising this. Given your specific request here, what do you think about changing the issue title to be a request to get linalg support on GPU? That's probably easier for us to add.
Done, thank you!
I want to know if  jax.numpy.linalg.inv ops can run on the GPU?
As it happens, jax.numpy.linalg.inv should run on GPU right now! (It might not be terribly fast, since it's using a QR decomposition instead of an LU decomposition, and one that isn't necessarily that well tuned.)
Another thing missing is `np.linalg.solve':

```
solve = np.linalg.solve(np.eye(3), np.ones(3))
Traceback (most recent call last):
  File "/opt/intellij-ue-2019.1/plugins/python/helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/numpy/linalg.py", line 237, in solve
    lu, pivots = lax_linalg.lu(a)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/lax_linalg.py", line 53, in lu
    lu, pivots = lu_p.bind(x)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/core.py", line 136, in bind
    return self.impl(*args, **kwargs)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/lax_linalg.py", line 358, in lu_impl
    lu, pivot = xla.apply_primitive(lu_p, operand)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/interpreters/xla.py", line 52, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/interpreters/xla.py", line 58, in xla_primitive_callable
    built_c = primitive_computation(prim, *shapes, **params)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/interpreters/xla.py", line 69, in primitive_computation
    xla_result = translation_rule(prim)(c, *xla_args, **params)
  File "/usr/local/google/_blaze_romann/513cef43ffae8d7478c0c7058e6a84e4/execroot/google3/blaze-out/k8-cuda9-py3-opt/bin/experimental/users/romann/ntk_tuner/train_and_eval.runfiles/google3/third_party/py/jax/lax_linalg.py", line 363, in lu_translation_rule
    "LU decomposition is only implemented on the CPU backend")
NotImplementedError: LU decomposition is only implemented on the CPU backend
```
There is now an LU decomposition implementation that works on GPU. However, it may not be the most performant (it's implemented in JAX itself). We still would do well to add a specialized GPU implementation that calls cuSolver or MAGMA.
(FYI, slight overlap with https://github.com/google/jax/issues/723)

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -16,6 +16,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from functools import partial
 import numpy as onp
 
 from jax.numpy import lax_numpy as np
@@ -34,6 +35,7 @@
 from jax.lax import (standard_primitive, standard_unop, binop_dtype_rule,
                      _float, _complex, _input_dtype, _broadcasting_select)
 from jax.lib import lapack
+from jax.lib import cusolver
 
 # traceables
 
@@ -81,6 +83,11 @@ def _T(x): return np.swapaxes(x, -1, -2)
 def _H(x): return np.conj(_T(x))
 def symmetrize(x): return (x + _H(x)) / 2
 
+def _unpack_tuple(f, n):
+  def g(c, *args, **kwargs):
+    t = f(c, *args, **kwargs)
+    return (c.GetTupleElement(t, i) for i in range(n))
+  return g
 
 # primitives
 
@@ -123,13 +130,18 @@ def _nan_like(c, operand):
     nan = c.Constant(onp.array(onp.nan, dtype=dtype))
   return c.Broadcast(nan, shape.dimensions())
 
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if hasattr(lapack, "potrf"):
+  _cpu_potrf = lapack.potrf
+else:
+  _cpu_potrf = _unpack_tuple(lapack.jax_potrf, 2)
+
 def cholesky_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
   dtype = shape.element_type().type
   if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
-    potrf_output = lapack.jax_potrf(c, operand, lower=True)
-    result = c.GetTupleElement(potrf_output, 0)
-    info = c.GetTupleElement(potrf_output, 1)
+    result, info = _cpu_potrf(c, operand, lower=True)
     return c.Select(c.Eq(info, c.ConstantS32Scalar(0)), result,
                     _nan_like(c, result))
   else:
@@ -163,15 +175,18 @@ def eig_abstract_eval(operand):
     w = vl = vr = operand
   return core.AbstractTuple((w, vl, vr))
 
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if hasattr(lapack, "geev"):
+  _cpu_geev = lapack.geev
+else:
+  _cpu_geev = _unpack_tuple(lapack.jax_geev, 4)
 
 def eig_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
   batch_dims = shape.dimensions()[:-2]
-  geev_out = lapack.jax_geev(c, operand)
-  w = c.GetTupleElement(geev_out, 0)
-  vl = c.GetTupleElement(geev_out, 1)
-  vr = c.GetTupleElement(geev_out, 2)
-  ok = c.Eq(c.GetTupleElement(geev_out, 3), c.ConstantS32Scalar(0))
+  w, vl, vr, info = _cpu_geev(c, operand)
+  ok = c.Eq(info, c.ConstantS32Scalar(0))
   w = _broadcasting_select(c, c.Reshape(ok, None, batch_dims + (1,)), w,
                            _nan_like(c, w))
   vl = _broadcasting_select(c, c.Reshape(ok, None, batch_dims + (1, 1)), vl,
@@ -219,13 +234,11 @@ def eigh_abstract_eval(operand, lower):
     v, w = operand, operand
   return core.AbstractTuple((v, w))
 
-def eigh_cpu_translation_rule(c, operand, lower):
+def _eigh_cpu_gpu_translation_rule(syevd_impl, c, operand, lower):
   shape = c.GetShape(operand)
   batch_dims = shape.dimensions()[:-2]
-  syevd_out = lapack.jax_syevd(c, operand, lower=lower)
-  v = c.GetTupleElement(syevd_out, 0)
-  w = c.GetTupleElement(syevd_out, 1)
-  ok = c.Eq(c.GetTupleElement(syevd_out, 2), c.ConstantS32Scalar(0))
+  v, w, info = syevd_impl(c, operand, lower=lower)
+  ok = c.Eq(info, c.ConstantS32Scalar(0))
   v = _broadcasting_select(c, c.Reshape(ok, None, batch_dims + (1, 1)), v,
                            _nan_like(c, v))
   w = _broadcasting_select(c, c.Reshape(ok, None, batch_dims + (1,)), w,
@@ -267,7 +280,22 @@ def eigh_batching_rule(batched_args, batch_dims, lower):
 eigh_p.def_abstract_eval(eigh_abstract_eval)
 xla.translations[eigh_p] = eigh_translation_rule
 ad.primitive_jvps[eigh_p] = eigh_jvp_rule
-xla.backend_specific_translations['cpu'][eigh_p] = eigh_cpu_translation_rule
+
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if hasattr(lapack, "syevd"):
+  _cpu_syevd = lapack.syevd
+else:
+  _cpu_syevd = _unpack_tuple(lapack.jax_syevd, 3)
+
+xla.backend_specific_translations['cpu'][eigh_p] = partial(
+  _eigh_cpu_gpu_translation_rule, _cpu_syevd)
+
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if cusolver:
+  xla.backend_specific_translations['gpu'][eigh_p] = partial(
+    _eigh_cpu_gpu_translation_rule, cusolver.syevd)
 batching.primitive_batchers[eigh_p] = eigh_batching_rule
 
 
@@ -522,14 +550,13 @@ def _lu_batching_rule(batched_args, batch_dims):
   x = batching.bdim_at_front(x, bd)
   return lu_p.bind(x), 0
 
-def _lu_cpu_translation_rule(c, operand):
+def _lu_cpu_gpu_translation_rule(getrf_impl, c, operand):
   shape = c.GetShape(operand)
   batch_dims = shape.dimensions()[:-2]
-  getrf_out = lapack.jax_getrf(c, operand)
-  lu = c.GetTupleElement(getrf_out, 0)
+  lu, pivot, info = getrf_impl(c, operand)
   # Subtract 1 from the pivot to get 0-based indices.
-  pivot = c.Sub(c.GetTupleElement(getrf_out, 1), c.ConstantS32Scalar(1))
-  ok = c.Eq(c.GetTupleElement(getrf_out, 2), c.ConstantS32Scalar(0))
+  pivot = c.Sub(pivot, c.ConstantS32Scalar(1))
+  ok = c.Eq(info, c.ConstantS32Scalar(0))
   lu = _broadcasting_select(c, c.Reshape(ok, None, batch_dims + (1, 1)), lu,
                             _nan_like(c, lu))
   return c.Tuple(lu, pivot)
@@ -541,7 +568,20 @@ def _lu_cpu_translation_rule(c, operand):
 xla.translations[lu_p] = xla.lower_fun(_lu_python, instantiate=True)
 ad.primitive_jvps[lu_p] = _lu_jvp_rule
 batching.primitive_batchers[lu_p] = _lu_batching_rule
-xla.backend_specific_translations['cpu'][lu_p] = _lu_cpu_translation_rule
+
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if hasattr(lapack, "getrf"):
+  _cpu_getrf = lapack.getrf
+else:
+  _cpu_getrf = _unpack_tuple(lapack.jax_getrf, 3)
+
+xla.backend_specific_translations['cpu'][lu_p] = partial(
+  _lu_cpu_gpu_translation_rule, _cpu_getrf)
+
+if cusolver:
+  xla.backend_specific_translations['gpu'][lu_p] = partial(
+    _lu_cpu_gpu_translation_rule, cusolver.getrf)
 
 
 def lu_pivots_to_permutation(swaps, m):
@@ -681,16 +721,13 @@ def svd_jvp_rule(primals, tangents, full_matrices, compute_uv):
     dV = dV + np.dot(np.eye(n) - np.dot(V, Vt), np.dot(np.conj(dA).T, U)) / s_dim
   return core.pack((s, U, Vt)), core.pack((ds, dU, dV.T))
 
-def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
+def _svd_cpu_gpu_translation_rule(gesvd_impl, c, operand, full_matrices, compute_uv):
   shape = c.GetShape(operand)
   dtype = shape.element_type().type
   if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
-    gesdd_out = lapack.jax_gesdd(c, operand, full_matrices=full_matrices,
-                                 compute_uv=compute_uv)
-    s = c.GetTupleElement(gesdd_out, 0)
-    u = c.GetTupleElement(gesdd_out, 1)
-    vt = c.GetTupleElement(gesdd_out, 2)
-    ok = c.Eq(c.GetTupleElement(gesdd_out, 3), c.ConstantS32Scalar(0))
+    s, u, vt, info = gesvd_impl(c, operand, full_matrices=full_matrices,
+                                compute_uv=compute_uv)
+    ok = c.Eq(info, c.ConstantS32Scalar(0))
     s = _broadcasting_select(c, c.Reshape(ok, None, (1,)), s,
                              _nan_like(c, s))
     u = _broadcasting_select(c, c.Reshape(ok, None, (1, 1)), u,
@@ -711,7 +748,22 @@ def svd_batching_rule(batched_args, batch_dims, full_matrices, compute_uv):
 svd_p = Primitive('svd')
 svd_p.def_impl(svd_impl)
 svd_p.def_abstract_eval(svd_abstract_eval)
-xla.translations[svd_p] = svd_translation_rule
-xla.backend_specific_translations['cpu'][svd_p] = svd_cpu_translation_rule
 ad.primitive_jvps[svd_p] = svd_jvp_rule
 batching.primitive_batchers[svd_p] = svd_batching_rule
+xla.translations[svd_p] = svd_translation_rule
+
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if hasattr(lapack, "gesdd"):
+  _cpu_gesdd = lapack.gesdd
+else:
+  _cpu_gesdd = _unpack_tuple(lapack.jax_gesdd, 4)
+
+xla.backend_specific_translations['cpu'][svd_p] = partial(
+  _svd_cpu_gpu_translation_rule, _cpu_gesdd)
+
+# TODO(phawkins): remove if-condition after increasing minimum Jaxlib version to
+# 0.1.23.
+if cusolver:
+  xla.backend_specific_translations['gpu'][svd_p] = partial(
+    _svd_cpu_gpu_translation_rule, cusolver.gesvd)
diff --git a/jax/lib/__init__.py b/jax/lib/__init__.py
--- a/jax/lib/__init__.py
+++ b/jax/lib/__init__.py
@@ -47,3 +47,10 @@ def _check_jaxlib_version():
   from jaxlib import pytree
 except ImportError:
   pytree = None
+
+# TODO(phawkins): make the import unconditional when the minimum Jaxlib version
+# has been increased to 0.1.23.
+try:
+  from jaxlib import cusolver
+except ImportError:
+  cusolver = None
diff --git a/jaxlib/cusolver.py b/jaxlib/cusolver.py
new file mode 100644
--- /dev/null
+++ b/jaxlib/cusolver.py
@@ -0,0 +1,182 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from jaxlib import xla_client
+
+try:
+  from jaxlib import cusolver_kernels
+  for _name, _value in cusolver_kernels.registrations().items():
+    xla_client.register_custom_call_target(_name, _value, platform="gpu")
+except ImportError:
+  pass
+
+_Shape = xla_client.Shape
+
+
+def _real_type(dtype):
+  """Returns the real equivalent of 'dtype'."""
+  if dtype == np.float32:
+    return np.float32
+  elif dtype == np.float64:
+    return np.float64
+  elif dtype == np.complex64:
+    return np.float32
+  elif dtype == np.complex128:
+    return np.float64
+  else:
+    raise NotImplementedError("Unsupported dtype {}".format(dtype))
+
+
+def getrf(c, a):
+  """LU decomposition."""
+  a_shape = c.GetShape(a)
+  dtype = a_shape.element_type()
+  dims = a_shape.dimensions()
+  assert len(dims) >= 2
+  m, n = dims[-2:]
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  b = 1
+  for d in batch_dims:
+    b *= d
+
+  lwork, opaque = cusolver_kernels.build_getrf_descriptor(
+      np.dtype(dtype), b, m, n)
+  out = c.CustomCall(
+      b"cusolver_getrf",
+      operands=(a,),
+      shape_with_layout=_Shape.tuple_shape((
+          _Shape.array_shape(
+              dtype, batch_dims + (m, n),
+              (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),
+          _Shape.array_shape(dtype, (lwork,), (0,)),
+          _Shape.array_shape(
+              np.dtype(np.int32), batch_dims + (min(m, n),),
+              tuple(range(num_bd, -1, -1))),
+          _Shape.array_shape(
+              np.dtype(np.int32), batch_dims, tuple(range(num_bd - 1, -1, -1))),
+      )),
+      operand_shapes_with_layout=(_Shape.array_shape(
+          dtype, batch_dims + (m, n),
+          (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),),
+      opaque=opaque)
+  return (c.GetTupleElement(out, 0), c.GetTupleElement(out, 2),
+          c.GetTupleElement(out, 3))
+
+
+def syevd(c, a, lower=False):
+  """Symmetric (Hermitian) eigendecomposition."""
+
+  a_shape = c.GetShape(a)
+  dtype = a_shape.element_type()
+  dims = a_shape.dimensions()
+  assert len(dims) >= 2
+  m, n = dims[-2:]
+  assert m == n
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  b = 1
+  for d in batch_dims:
+    b *= d
+  layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
+
+  lwork, opaque = cusolver_kernels.build_syevd_descriptor(
+      np.dtype(dtype), lower, b, n)
+  eigvals_type = _real_type(dtype)
+
+  out = c.CustomCall(
+      b"cusolver_syevd",
+      operands=(a,),
+      shape_with_layout=_Shape.tuple_shape((
+          _Shape.array_shape(dtype, dims, layout),
+          _Shape.array_shape(
+              np.dtype(eigvals_type), batch_dims + (n,),
+              tuple(range(num_bd, -1, -1))),
+          _Shape.array_shape(
+              np.dtype(np.int32), batch_dims,
+              tuple(range(num_bd - 1, -1, -1))),
+          _Shape.array_shape(dtype, (lwork,), (0,))
+      )),
+      operand_shapes_with_layout=(
+          _Shape.array_shape(dtype, dims, layout),
+      ),
+      opaque=opaque)
+  return (c.GetTupleElement(out, 0), c.GetTupleElement(out, 1),
+          c.GetTupleElement(out, 2))
+
+
+def gesvd(c, a, full_matrices=True, compute_uv=True):
+  """Singular value decomposition."""
+
+  a_shape = c.GetShape(a)
+  dtype = a_shape.element_type()
+  b = 1
+  m, n = a_shape.dimensions()
+  singular_vals_dtype = _real_type(dtype)
+
+  if m < n:
+    lwork, opaque = cusolver_kernels.build_gesvd_descriptor(
+        np.dtype(dtype), b, n, m, compute_uv, full_matrices)
+    out = c.CustomCall(
+        b"cusolver_gesvd",
+        operands=(a,),
+        shape_with_layout=_Shape.tuple_shape((
+            _Shape.array_shape(dtype, (m, n), (1, 0)),
+            _Shape.array_shape(np.dtype(singular_vals_dtype), (min(m, n),), (0,)),
+            _Shape.array_shape(dtype, (n, n), (1, 0)),
+            _Shape.array_shape(dtype, (m, m), (1, 0)),
+            _Shape.array_shape(np.dtype(np.int32), (), ()),
+            _Shape.array_shape(dtype, (lwork,), (0,)),
+        )),
+        operand_shapes_with_layout=(
+            _Shape.array_shape(dtype, (m, n), (1, 0)),
+        ),
+        opaque=opaque)
+    s = c.GetTupleElement(out, 1)
+    vt = c.GetTupleElement(out, 2)
+    u = c.GetTupleElement(out, 3)
+    info = c.GetTupleElement(out, 4)
+  else:
+    lwork, opaque = cusolver_kernels.build_gesvd_descriptor(
+        np.dtype(dtype), b, m, n, compute_uv, full_matrices)
+
+    out = c.CustomCall(
+        b"cusolver_gesvd",
+        operands=(a,),
+        shape_with_layout=_Shape.tuple_shape((
+            _Shape.array_shape(dtype, (m, n), (0, 1)),
+            _Shape.array_shape(np.dtype(singular_vals_dtype), (min(m, n),), (0,)),
+            _Shape.array_shape(dtype, (m, m), (0, 1)),
+            _Shape.array_shape(dtype, (n, n), (0, 1)),
+            _Shape.array_shape(np.dtype(np.int32), (), ()),
+            _Shape.array_shape(dtype, (lwork,), (0,)),
+        )),
+        operand_shapes_with_layout=(
+            _Shape.array_shape(dtype, (m, n), (0, 1)),
+        ),
+        opaque=opaque)
+    s = c.GetTupleElement(out, 1)
+    u = c.GetTupleElement(out, 2)
+    vt = c.GetTupleElement(out, 3)
+    info = c.GetTupleElement(out, 4)
+  if not full_matrices:
+    u = c.Slice(u, (0, 0), (m, min(m, n)))
+    vt = c.Slice(vt, (0, 0), (min(m, n), n))
+  return s, u, vt, info


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -163,8 +163,8 @@ def testEigBatching(self, shape, dtype, rng):
       for lower in [False, True]
       for rng in [jtu.rand_default()]))
   # TODO(phawkins): enable when there is an eigendecomposition implementation
-  # for GPU/TPU.
-  @jtu.skip_on_devices("gpu", "tpu")
+  # for TPU.
+  @jtu.skip_on_devices("tpu")
   def testEigh(self, n, dtype, lower, rng):
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng((n, n), dtype)]
@@ -196,8 +196,8 @@ def norm(x):
       for rng in [jtu.rand_default()]
       for lower in [True, False]))
   # TODO(phawkins): enable when there is an eigendecomposition implementation
-  # for GPU/TPU.
-  @jtu.skip_on_devices("gpu", "tpu")
+  # for TPU.
+  @jtu.skip_on_devices("tpu")
   def testEighGrad(self, shape, dtype, rng, lower):
     self.skipTest("Test fails with numeric errors.")
     uplo = "L" if lower else "U"
@@ -224,8 +224,8 @@ def testEighGrad(self, shape, dtype, rng, lower):
       for lower in [True, False]
       for eps in [1e-4]))
   # TODO(phawkins): enable when there is an eigendecomposition implementation
-  # for GPU/TPU.
-  @jtu.skip_on_devices("gpu", "tpu")
+  # for TPU.
+  @jtu.skip_on_devices("tpu")
   def testEighGradVectorComplex(self, shape, dtype, rng, lower, eps):
     _skip_if_unsupported_type(dtype)
     # Special case to test for complex eigenvector grad correctness.
@@ -263,7 +263,7 @@ def testEighGradVectorComplex(self, shape, dtype, rng, lower, eps):
       for shape in [(1, 1), (4, 4), (5, 5)]
       for dtype in float_types + complex_types
       for rng in [jtu.rand_default()]))
-  @jtu.skip_on_devices("gpu", "tpu")
+  @jtu.skip_on_devices("tpu")
   def testEighBatching(self, shape, dtype, rng):
     _skip_if_unsupported_type(dtype)
     shape = (10,) + shape
@@ -318,7 +318,7 @@ def testNorm(self, shape, dtype, ord, axis, keepdims, rng):
       for full_matrices in [False, True]
       for compute_uv in [False, True]
       for rng in [jtu.rand_default()]))
-  @jtu.skip_on_devices("gpu", "tpu")
+  @jtu.skip_on_devices("tpu")
   def testSVD(self, m, n, dtype, full_matrices, compute_uv, rng):
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng((m, n), dtype)]
@@ -414,7 +414,7 @@ def compare_orthogonal(q1, q2):
     if not full_matrices and m >= n:
         jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))
 
-  @jtu.skip_on_devices("gpu", "tpu")
+  @jtu.skip_on_devices("tpu")
   def testQrBatching(self):
     shape = (10, 4, 5)
     dtype = np.float32
@@ -476,7 +476,7 @@ def args_maker():
     self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)
 
   # Regression test for incorrect type for eigenvalues of a complex matrix.
-  @jtu.skip_on_devices("gpu", "tpu")
+  @jtu.skip_on_devices("tpu")
   def testIssue669(self):
     def test(x):
       val, vec = np.linalg.eigh(x)
@@ -499,9 +499,9 @@ class ScipyLinalgTest(jtu.JaxTestCase):
   def testLu(self, shape, dtype, rng):
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng(shape, dtype)]
-
-    self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
-                            check_dtypes=True, tol=1e-3)
+    x, = args_maker()
+    p, l, u = jsp.linalg.lu(x)
+    self.assertAllClose(x, onp.matmul(p, onp.matmul(l, u)), check_dtypes=True)
     self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)
 
   # TODO(phawkins): figure out why this test fails on Travis and reenable.
@@ -555,8 +555,13 @@ def testLuFactor(self, n, dtype, rng):
     _skip_if_unsupported_type(dtype)
     args_maker = lambda: [rng((n, n), dtype)]
 
-    self._CheckAgainstNumpy(jsp.linalg.lu_factor, osp.linalg.lu_factor,
-                            args_maker, check_dtypes=True, tol=1e-3)
+    x, = args_maker()
+    lu, piv = jsp.linalg.lu_factor(x)
+    l = onp.tril(lu, -1) + onp.eye(n, dtype=dtype)
+    u = onp.triu(lu)
+    for i in range(n):
+      x[[i, piv[i]],] = x[[piv[i], i],]
+    self.assertAllClose(x, onp.matmul(l, u), check_dtypes=True, rtol=1e-3)
     self._CompileAndCheck(jsp.linalg.lu_factor, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(