# Copyright 2018 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.

"""Tests for the LAPAX linear algebra module."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from functools import partial
import itertools

import numpy as onp
import scipy as osp

from absl.testing import absltest
from absl.testing import parameterized

from jax import jvp
from jax import numpy as np
from jax import scipy as jsp
from jax import test_util as jtu
from jax.lib import xla_bridge

from jaxlib import lapack

from jax.config import config
config.parse_flags_with_absl()

T = lambda x: onp.swapaxes(x, -1, -2)


def float_types():
  return set(onp.dtype(xla_bridge.canonicalize_dtype(dtype))
             for dtype in [onp.float32, onp.float64])

def complex_types():
  return {onp.complex64}


class NumpyLinalgTest(jtu.JaxTestCase):

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testCholesky(self, shape, dtype, rng):
    def args_maker():
      a = rng(shape, dtype)
      return [onp.matmul(a, np.conj(T(a)))]

    self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [0, 4, 5, 50]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testDet(self, n, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng((n, n), dtype)]

    self._CheckAgainstNumpy(onp.linalg.det, np.linalg.det, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.det, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [0, 4, 10, 200]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testSlogdet(self, n, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng((n, n), dtype)]

    self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_fullmatrices={}".format(
          jtu.format_shape_dtype_string(shape, dtype), full_matrices),
       "shape": shape, "dtype": dtype, "full_matrices": full_matrices,
       "rng": rng}
      for shape in [(1, 1), (3, 4), (2, 10, 5), (2, 200, 100)]
      for dtype in float_types()
      for full_matrices in [False, True]
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("cpu")
  def testQr(self, shape, dtype, full_matrices, rng):
    m, n = shape[-2:]

    if full_matrices:
      mode, k = "complete", m
    else:
      mode, k = "reduced", min(m, n)

    a = rng(shape, dtype)
    lq, lr = np.linalg.qr(a, mode=mode)

    # onp.linalg.qr doesn't support broadcasting. But it seems like an
    # inevitable extension so we support it in our version.
    nq = onp.zeros(shape[:-2] + (m, k), dtype)
    nr = onp.zeros(shape[:-2] + (k, n), dtype)
    for index in onp.ndindex(*shape[:-2]):
      nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)

    max_rank = max(m, n)

    # Norm, adjusted for dimension and type.
    def norm(x):
      n = onp.linalg.norm(x, axis=(-2, -1))
      return n / (max_rank * onp.finfo(dtype).eps)

    def compare_orthogonal(q1, q2):
      # Q is unique up to sign, so normalize the sign first.
      sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
      phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
      q1 *= phases
      self.assertTrue(onp.all(norm(q1 - q2) < 30))

    # Check a ~= qr
    self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))

    # Compare the first 'k' vectors of Q; the remainder form an arbitrary
    # orthonormal basis for the null space.
    compare_orthogonal(nq[..., :k], lq[..., :k])

    # Check that q is close to unitary.
    self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))

    if not full_matrices and m >= n:
        jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype)),
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lhs_shape, rhs_shape in [
          ((1, 1), (1, 1)),
          ((4, 4), (4,)),
          ((8, 8), (8, 4)),
      ]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testSolve(self, lhs_shape, rhs_shape, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]

    self._CheckAgainstNumpy(onp.linalg.solve, np.linalg.solve, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.solve, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (5, 5, 5)]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testInv(self, shape, dtype, rng):
    def args_maker():
      invertible = False
      while not invertible:
        a = rng(shape, dtype)
        try:
          onp.linalg.inv(a)
          invertible = True
        except onp.linalg.LinAlgError:
          pass
      return [a]

    self._CheckAgainstNumpy(onp.linalg.inv, np.linalg.inv, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)


class ScipyLinalgTest(jtu.JaxTestCase):

  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 5), (10, 5), (50, 50)]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testLu(self, shape, dtype, rng):
    # TODO(phawkins): remove this after a jaxlib release.
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng(shape, dtype)]

    self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [1, 4, 5, 200]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testLuFactor(self, n, dtype, rng):
    # TODO(phawkins): remove this after a jaxlib release.
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    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)
    self._CompileAndCheck(jsp.linalg.lu_factor, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_sym_pos={}_lower={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           sym_pos, lower),
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "sym_pos": sym_pos, "lower": lower, "rng": rng}
      for lhs_shape, rhs_shape in [
          ((1, 1), (1, 1)),
          ((4, 4), (4,)),
          ((8, 8), (8, 4)),
      ]
      for sym_pos, lower in [
        (False, False),
        (True, False),
        (True, True),
      ]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=sym_pos, lower=lower)
    jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=sym_pos, lower=lower)

    def args_maker():
      a = rng(lhs_shape, dtype)
      if sym_pos:
        a = onp.matmul(a, onp.conj(T(a)))
        a = onp.tril(a) if lower else onp.triu(a)
      return [a, rng(rhs_shape, dtype)]

    self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_lower={}_transposea={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           lower, transpose_a),
       "lower": lower, "transpose_a": transpose_a,
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lower, transpose_a in itertools.product([False, True], repeat=2)
      for lhs_shape, rhs_shape in [
          ((4, 4), (4,)),
          ((4, 4), (4, 3)),
          ((2, 8, 8), (2, 8, 10)),
      ]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testSolveTriangular(self, lower, transpose_a, lhs_shape,
                                 rhs_shape, dtype, rng):
    k = rng(lhs_shape, dtype)
    l = onp.linalg.cholesky(onp.matmul(k, T(k))
                            + lhs_shape[-1] * onp.eye(lhs_shape[-1]))
    l = l.astype(k.dtype)
    b = rng(rhs_shape, dtype)

    a = l if lower else T(l)
    inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
    if len(lhs_shape) == len(rhs_shape):
      onp_ans = onp.matmul(inv, b)
    else:
      onp_ans = onp.einsum("...ij,...j->...i", inv, b)

    # The standard scipy.linalg.solve_triangular doesn't support broadcasting.
    # But it seems like an inevitable extension so we support it.
    ans = jsp.linalg.solve_triangular(
        l if lower else T(l), b, trans=1 if transpose_a else 0, lower=lower)

    self.assertAllClose(onp_ans, ans, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_lower={}_transposea={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           lower, transpose_a),
       "lower": lower, "transpose_a": transpose_a,
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lower, transpose_a in itertools.product([False, True], repeat=2)
      for lhs_shape, rhs_shape in [
          ((4, 4), (4,)),
          ((4, 4), (4, 3)),
          ((2, 8, 8), (2, 8, 10)),
      ]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
                                     rhs_shape, dtype, rng):
    # TODO(frostig): change ensemble to support a bigger rtol
    self.skipTest("rtol does not cover all devices and precision modes")
    A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1], dtype=dtype))
    A = A if lower else T(A)
    B = rng(rhs_shape, dtype)
    f = partial(jsp.linalg.solve_triangular, lower=lower,
                trans=1 if transpose_a else 0)
    jtu.check_grads(f, (A, B), 2, rtol=1e-3)

if __name__ == "__main__":
  absltest.main()

68df557696da12322983c7d5ffc8842a6dae4155
<mask token>


class NumpyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)] for dtype in
        float_types() for rng in [jtu.rand_default()]))
    def testCholesky(self, shape, dtype, rng):

        def args_maker():
            a = rng(shape, dtype)
            return [onp.matmul(a, np.conj(T(a)))]
        self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True
            )
    <mask token>

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 10, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSlogdet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}_fullmatrices={}'.format(jtu.format_shape_dtype_string(
        shape, dtype), full_matrices), 'shape': shape, 'dtype': dtype,
        'full_matrices': full_matrices, 'rng': rng} for shape in [(1, 1), (
        3, 4), (2, 10, 5), (2, 200, 100)] for dtype in float_types() for
        full_matrices in [False, True] for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('cpu')
    def testQr(self, shape, dtype, full_matrices, rng):
        m, n = shape[-2:]
        if full_matrices:
            mode, k = 'complete', m
        else:
            mode, k = 'reduced', min(m, n)
        a = rng(shape, dtype)
        lq, lr = np.linalg.qr(a, mode=mode)
        nq = onp.zeros(shape[:-2] + (m, k), dtype)
        nr = onp.zeros(shape[:-2] + (k, n), dtype)
        for index in onp.ndindex(*shape[:-2]):
            nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)
        max_rank = max(m, n)

        def norm(x):
            n = onp.linalg.norm(x, axis=(-2, -1))
            return n / (max_rank * onp.finfo(dtype).eps)

        def compare_orthogonal(q1, q2):
            sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
            phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
            q1 *= phases
            self.assertTrue(onp.all(norm(q1 - q2) < 30))
        self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))
        compare_orthogonal(nq[..., :k], lq[..., :k])
        self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))
        if not full_matrices and m >= n:
            jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))
    <mask token>
    <mask token>


class ScipyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 5), (10, 5), (50, 50)] for dtype in float_types() |
        complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLu(self, shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(shape, dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [1, 4, 5, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLuFactor(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu_factor, osp.linalg.lu_factor,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu_factor, args_maker,
            check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_sym_pos={}_lower={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), sym_pos, lower),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'sym_pos': sym_pos, 'lower': lower, 'rng': rng} for lhs_shape,
        rhs_shape in [((1, 1), (1, 1)), ((4, 4), (4,)), ((8, 8), (8, 4))] for
        sym_pos, lower in [(False, False), (True, False), (True, True)] for
        dtype in float_types() | complex_types() for rng in [jtu.
        rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)
        jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)

        def args_maker():
            a = rng(lhs_shape, dtype)
            if sym_pos:
                a = onp.matmul(a, onp.conj(T(a)))
                a = onp.tril(a) if lower else onp.triu(a)
            return [a, rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=
            True, tol=0.001)
        self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangular(self, lower, transpose_a, lhs_shape, rhs_shape,
        dtype, rng):
        k = rng(lhs_shape, dtype)
        l = onp.linalg.cholesky(onp.matmul(k, T(k)) + lhs_shape[-1] * onp.
            eye(lhs_shape[-1]))
        l = l.astype(k.dtype)
        b = rng(rhs_shape, dtype)
        a = l if lower else T(l)
        inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
        if len(lhs_shape) == len(rhs_shape):
            onp_ans = onp.matmul(inv, b)
        else:
            onp_ans = onp.einsum('...ij,...j->...i', inv, b)
        ans = jsp.linalg.solve_triangular(l if lower else T(l), b, trans=1 if
            transpose_a else 0, lower=lower)
        self.assertAllClose(onp_ans, ans, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
        rhs_shape, dtype, rng):
        self.skipTest('rtol does not cover all devices and precision modes')
        A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1],
            dtype=dtype))
        A = A if lower else T(A)
        B = rng(rhs_shape, dtype)
        f = partial(jsp.linalg.solve_triangular, lower=lower, trans=1 if
            transpose_a else 0)
        jtu.check_grads(f, (A, B), 2, rtol=0.001)


<mask token>

<mask token>


class NumpyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)] for dtype in
        float_types() for rng in [jtu.rand_default()]))
    def testCholesky(self, shape, dtype, rng):

        def args_maker():
            a = rng(shape, dtype)
            return [onp.matmul(a, np.conj(T(a)))]
        self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True
            )
    <mask token>

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 10, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSlogdet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}_fullmatrices={}'.format(jtu.format_shape_dtype_string(
        shape, dtype), full_matrices), 'shape': shape, 'dtype': dtype,
        'full_matrices': full_matrices, 'rng': rng} for shape in [(1, 1), (
        3, 4), (2, 10, 5), (2, 200, 100)] for dtype in float_types() for
        full_matrices in [False, True] for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('cpu')
    def testQr(self, shape, dtype, full_matrices, rng):
        m, n = shape[-2:]
        if full_matrices:
            mode, k = 'complete', m
        else:
            mode, k = 'reduced', min(m, n)
        a = rng(shape, dtype)
        lq, lr = np.linalg.qr(a, mode=mode)
        nq = onp.zeros(shape[:-2] + (m, k), dtype)
        nr = onp.zeros(shape[:-2] + (k, n), dtype)
        for index in onp.ndindex(*shape[:-2]):
            nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)
        max_rank = max(m, n)

        def norm(x):
            n = onp.linalg.norm(x, axis=(-2, -1))
            return n / (max_rank * onp.finfo(dtype).eps)

        def compare_orthogonal(q1, q2):
            sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
            phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
            q1 *= phases
            self.assertTrue(onp.all(norm(q1 - q2) < 30))
        self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))
        compare_orthogonal(nq[..., :k], lq[..., :k])
        self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))
        if not full_matrices and m >= n:
            jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))
    <mask token>

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (5, 5, 5)] for dtype in float_types() for
        rng in [jtu.rand_default()]))
    def testInv(self, shape, dtype, rng):

        def args_maker():
            invertible = False
            while not invertible:
                a = rng(shape, dtype)
                try:
                    onp.linalg.inv(a)
                    invertible = True
                except onp.linalg.LinAlgError:
                    pass
            return [a]
        self._CheckAgainstNumpy(onp.linalg.inv, np.linalg.inv, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)


class ScipyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 5), (10, 5), (50, 50)] for dtype in float_types() |
        complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLu(self, shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(shape, dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [1, 4, 5, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLuFactor(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu_factor, osp.linalg.lu_factor,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu_factor, args_maker,
            check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_sym_pos={}_lower={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), sym_pos, lower),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'sym_pos': sym_pos, 'lower': lower, 'rng': rng} for lhs_shape,
        rhs_shape in [((1, 1), (1, 1)), ((4, 4), (4,)), ((8, 8), (8, 4))] for
        sym_pos, lower in [(False, False), (True, False), (True, True)] for
        dtype in float_types() | complex_types() for rng in [jtu.
        rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)
        jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)

        def args_maker():
            a = rng(lhs_shape, dtype)
            if sym_pos:
                a = onp.matmul(a, onp.conj(T(a)))
                a = onp.tril(a) if lower else onp.triu(a)
            return [a, rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=
            True, tol=0.001)
        self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangular(self, lower, transpose_a, lhs_shape, rhs_shape,
        dtype, rng):
        k = rng(lhs_shape, dtype)
        l = onp.linalg.cholesky(onp.matmul(k, T(k)) + lhs_shape[-1] * onp.
            eye(lhs_shape[-1]))
        l = l.astype(k.dtype)
        b = rng(rhs_shape, dtype)
        a = l if lower else T(l)
        inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
        if len(lhs_shape) == len(rhs_shape):
            onp_ans = onp.matmul(inv, b)
        else:
            onp_ans = onp.einsum('...ij,...j->...i', inv, b)
        ans = jsp.linalg.solve_triangular(l if lower else T(l), b, trans=1 if
            transpose_a else 0, lower=lower)
        self.assertAllClose(onp_ans, ans, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
        rhs_shape, dtype, rng):
        self.skipTest('rtol does not cover all devices and precision modes')
        A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1],
            dtype=dtype))
        A = A if lower else T(A)
        B = rng(rhs_shape, dtype)
        f = partial(jsp.linalg.solve_triangular, lower=lower, trans=1 if
            transpose_a else 0)
        jtu.check_grads(f, (A, B), 2, rtol=0.001)


<mask token>

<mask token>


class NumpyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)] for dtype in
        float_types() for rng in [jtu.rand_default()]))
    def testCholesky(self, shape, dtype, rng):

        def args_maker():
            a = rng(shape, dtype)
            return [onp.matmul(a, np.conj(T(a)))]
        self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True
            )

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 5, 50] for dtype in 
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testDet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.det, np.linalg.det, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.det, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 10, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSlogdet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}_fullmatrices={}'.format(jtu.format_shape_dtype_string(
        shape, dtype), full_matrices), 'shape': shape, 'dtype': dtype,
        'full_matrices': full_matrices, 'rng': rng} for shape in [(1, 1), (
        3, 4), (2, 10, 5), (2, 200, 100)] for dtype in float_types() for
        full_matrices in [False, True] for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('cpu')
    def testQr(self, shape, dtype, full_matrices, rng):
        m, n = shape[-2:]
        if full_matrices:
            mode, k = 'complete', m
        else:
            mode, k = 'reduced', min(m, n)
        a = rng(shape, dtype)
        lq, lr = np.linalg.qr(a, mode=mode)
        nq = onp.zeros(shape[:-2] + (m, k), dtype)
        nr = onp.zeros(shape[:-2] + (k, n), dtype)
        for index in onp.ndindex(*shape[:-2]):
            nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)
        max_rank = max(m, n)

        def norm(x):
            n = onp.linalg.norm(x, axis=(-2, -1))
            return n / (max_rank * onp.finfo(dtype).eps)

        def compare_orthogonal(q1, q2):
            sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
            phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
            q1 *= phases
            self.assertTrue(onp.all(norm(q1 - q2) < 30))
        self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))
        compare_orthogonal(nq[..., :k], lq[..., :k])
        self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))
        if not full_matrices and m >= n:
            jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}'.format(jtu.format_shape_dtype_string(lhs_shape,
        dtype), jtu.format_shape_dtype_string(rhs_shape, dtype)),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'rng': rng} for lhs_shape, rhs_shape in [((1, 1), (1, 1)), ((4, 4),
        (4,)), ((8, 8), (8, 4))] for dtype in float_types() | complex_types
        () for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(onp.linalg.solve, np.linalg.solve,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.solve, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (5, 5, 5)] for dtype in float_types() for
        rng in [jtu.rand_default()]))
    def testInv(self, shape, dtype, rng):

        def args_maker():
            invertible = False
            while not invertible:
                a = rng(shape, dtype)
                try:
                    onp.linalg.inv(a)
                    invertible = True
                except onp.linalg.LinAlgError:
                    pass
            return [a]
        self._CheckAgainstNumpy(onp.linalg.inv, np.linalg.inv, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)


class ScipyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 5), (10, 5), (50, 50)] for dtype in float_types() |
        complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLu(self, shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(shape, dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [1, 4, 5, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLuFactor(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu_factor, osp.linalg.lu_factor,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu_factor, args_maker,
            check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_sym_pos={}_lower={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), sym_pos, lower),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'sym_pos': sym_pos, 'lower': lower, 'rng': rng} for lhs_shape,
        rhs_shape in [((1, 1), (1, 1)), ((4, 4), (4,)), ((8, 8), (8, 4))] for
        sym_pos, lower in [(False, False), (True, False), (True, True)] for
        dtype in float_types() | complex_types() for rng in [jtu.
        rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)
        jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)

        def args_maker():
            a = rng(lhs_shape, dtype)
            if sym_pos:
                a = onp.matmul(a, onp.conj(T(a)))
                a = onp.tril(a) if lower else onp.triu(a)
            return [a, rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=
            True, tol=0.001)
        self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangular(self, lower, transpose_a, lhs_shape, rhs_shape,
        dtype, rng):
        k = rng(lhs_shape, dtype)
        l = onp.linalg.cholesky(onp.matmul(k, T(k)) + lhs_shape[-1] * onp.
            eye(lhs_shape[-1]))
        l = l.astype(k.dtype)
        b = rng(rhs_shape, dtype)
        a = l if lower else T(l)
        inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
        if len(lhs_shape) == len(rhs_shape):
            onp_ans = onp.matmul(inv, b)
        else:
            onp_ans = onp.einsum('...ij,...j->...i', inv, b)
        ans = jsp.linalg.solve_triangular(l if lower else T(l), b, trans=1 if
            transpose_a else 0, lower=lower)
        self.assertAllClose(onp_ans, ans, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
        rhs_shape, dtype, rng):
        self.skipTest('rtol does not cover all devices and precision modes')
        A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1],
            dtype=dtype))
        A = A if lower else T(A)
        B = rng(rhs_shape, dtype)
        f = partial(jsp.linalg.solve_triangular, lower=lower, trans=1 if
            transpose_a else 0)
        jtu.check_grads(f, (A, B), 2, rtol=0.001)


<mask token>

<mask token>


def complex_types():
    return {onp.complex64}


class NumpyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)] for dtype in
        float_types() for rng in [jtu.rand_default()]))
    def testCholesky(self, shape, dtype, rng):

        def args_maker():
            a = rng(shape, dtype)
            return [onp.matmul(a, np.conj(T(a)))]
        self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True
            )

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 5, 50] for dtype in 
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testDet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.det, np.linalg.det, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.det, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [0, 4, 10, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSlogdet(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}_fullmatrices={}'.format(jtu.format_shape_dtype_string(
        shape, dtype), full_matrices), 'shape': shape, 'dtype': dtype,
        'full_matrices': full_matrices, 'rng': rng} for shape in [(1, 1), (
        3, 4), (2, 10, 5), (2, 200, 100)] for dtype in float_types() for
        full_matrices in [False, True] for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('cpu')
    def testQr(self, shape, dtype, full_matrices, rng):
        m, n = shape[-2:]
        if full_matrices:
            mode, k = 'complete', m
        else:
            mode, k = 'reduced', min(m, n)
        a = rng(shape, dtype)
        lq, lr = np.linalg.qr(a, mode=mode)
        nq = onp.zeros(shape[:-2] + (m, k), dtype)
        nr = onp.zeros(shape[:-2] + (k, n), dtype)
        for index in onp.ndindex(*shape[:-2]):
            nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)
        max_rank = max(m, n)

        def norm(x):
            n = onp.linalg.norm(x, axis=(-2, -1))
            return n / (max_rank * onp.finfo(dtype).eps)

        def compare_orthogonal(q1, q2):
            sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
            phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
            q1 *= phases
            self.assertTrue(onp.all(norm(q1 - q2) < 30))
        self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))
        compare_orthogonal(nq[..., :k], lq[..., :k])
        self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))
        if not full_matrices and m >= n:
            jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}'.format(jtu.format_shape_dtype_string(lhs_shape,
        dtype), jtu.format_shape_dtype_string(rhs_shape, dtype)),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'rng': rng} for lhs_shape, rhs_shape in [((1, 1), (1, 1)), ((4, 4),
        (4,)), ((8, 8), (8, 4))] for dtype in float_types() | complex_types
        () for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(onp.linalg.solve, np.linalg.solve,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.solve, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 4), (2, 5, 5), (200, 200), (5, 5, 5)] for dtype in float_types() for
        rng in [jtu.rand_default()]))
    def testInv(self, shape, dtype, rng):

        def args_maker():
            invertible = False
            while not invertible:
                a = rng(shape, dtype)
                try:
                    onp.linalg.inv(a)
                    invertible = True
                except onp.linalg.LinAlgError:
                    pass
            return [a]
        self._CheckAgainstNumpy(onp.linalg.inv, np.linalg.inv, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)


class ScipyLinalgTest(jtu.JaxTestCase):

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_shape={}'.format(jtu.format_shape_dtype_string(shape, dtype)),
        'shape': shape, 'dtype': dtype, 'rng': rng} for shape in [(1, 1), (
        4, 5), (10, 5), (50, 50)] for dtype in float_types() |
        complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLu(self, shape, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng(shape, dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
            check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_n={}'.format(jtu.format_shape_dtype_string((n, n), dtype)), 'n':
        n, 'dtype': dtype, 'rng': rng} for n in [1, 4, 5, 200] for dtype in
        float_types() | complex_types() for rng in [jtu.rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testLuFactor(self, n, dtype, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        args_maker = lambda : [rng((n, n), dtype)]
        self._CheckAgainstNumpy(jsp.linalg.lu_factor, osp.linalg.lu_factor,
            args_maker, check_dtypes=True, tol=0.001)
        self._CompileAndCheck(jsp.linalg.lu_factor, args_maker,
            check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_sym_pos={}_lower={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), sym_pos, lower),
        'lhs_shape': lhs_shape, 'rhs_shape': rhs_shape, 'dtype': dtype,
        'sym_pos': sym_pos, 'lower': lower, 'rng': rng} for lhs_shape,
        rhs_shape in [((1, 1), (1, 1)), ((4, 4), (4,)), ((8, 8), (8, 4))] for
        sym_pos, lower in [(False, False), (True, False), (True, True)] for
        dtype in float_types() | complex_types() for rng in [jtu.
        rand_default()]))
    @jtu.skip_on_devices('gpu', 'tpu')
    def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
        if not hasattr(lapack, 'jax_getrf'):
            self.skipTest('No LU implementation available')
        osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)
        jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=
            sym_pos, lower=lower)

        def args_maker():
            a = rng(lhs_shape, dtype)
            if sym_pos:
                a = onp.matmul(a, onp.conj(T(a)))
                a = onp.tril(a) if lower else onp.triu(a)
            return [a, rng(rhs_shape, dtype)]
        self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker, check_dtypes=
            True, tol=0.001)
        self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangular(self, lower, transpose_a, lhs_shape, rhs_shape,
        dtype, rng):
        k = rng(lhs_shape, dtype)
        l = onp.linalg.cholesky(onp.matmul(k, T(k)) + lhs_shape[-1] * onp.
            eye(lhs_shape[-1]))
        l = l.astype(k.dtype)
        b = rng(rhs_shape, dtype)
        a = l if lower else T(l)
        inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
        if len(lhs_shape) == len(rhs_shape):
            onp_ans = onp.matmul(inv, b)
        else:
            onp_ans = onp.einsum('...ij,...j->...i', inv, b)
        ans = jsp.linalg.solve_triangular(l if lower else T(l), b, trans=1 if
            transpose_a else 0, lower=lower)
        self.assertAllClose(onp_ans, ans, check_dtypes=True)

    @parameterized.named_parameters(jtu.cases_from_list({'testcase_name':
        '_lhs={}_rhs={}_lower={}_transposea={}'.format(jtu.
        format_shape_dtype_string(lhs_shape, dtype), jtu.
        format_shape_dtype_string(rhs_shape, dtype), lower, transpose_a),
        'lower': lower, 'transpose_a': transpose_a, 'lhs_shape': lhs_shape,
        'rhs_shape': rhs_shape, 'dtype': dtype, 'rng': rng} for lower,
        transpose_a in itertools.product([False, True], repeat=2) for 
        lhs_shape, rhs_shape in [((4, 4), (4,)), ((4, 4), (4, 3)), ((2, 8, 
        8), (2, 8, 10))] for dtype in float_types() for rng in [jtu.
        rand_default()]))
    def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
        rhs_shape, dtype, rng):
        self.skipTest('rtol does not cover all devices and precision modes')
        A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1],
            dtype=dtype))
        A = A if lower else T(A)
        B = rng(rhs_shape, dtype)
        f = partial(jsp.linalg.solve_triangular, lower=lower, trans=1 if
            transpose_a else 0)
        jtu.check_grads(f, (A, B), 2, rtol=0.001)


<mask token>

# Copyright 2018 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.

"""Tests for the LAPAX linear algebra module."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from functools import partial
import itertools

import numpy as onp
import scipy as osp

from absl.testing import absltest
from absl.testing import parameterized

from jax import jvp
from jax import numpy as np
from jax import scipy as jsp
from jax import test_util as jtu
from jax.lib import xla_bridge

from jaxlib import lapack

from jax.config import config
config.parse_flags_with_absl()

T = lambda x: onp.swapaxes(x, -1, -2)


def float_types():
  return set(onp.dtype(xla_bridge.canonicalize_dtype(dtype))
             for dtype in [onp.float32, onp.float64])

def complex_types():
  return {onp.complex64}


class NumpyLinalgTest(jtu.JaxTestCase):

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (1000, 0, 0)]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testCholesky(self, shape, dtype, rng):
    def args_maker():
      a = rng(shape, dtype)
      return [onp.matmul(a, np.conj(T(a)))]

    self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [0, 4, 5, 50]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testDet(self, n, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng((n, n), dtype)]

    self._CheckAgainstNumpy(onp.linalg.det, np.linalg.det, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.det, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [0, 4, 10, 200]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testSlogdet(self, n, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng((n, n), dtype)]

    self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name": "_shape={}_fullmatrices={}".format(
          jtu.format_shape_dtype_string(shape, dtype), full_matrices),
       "shape": shape, "dtype": dtype, "full_matrices": full_matrices,
       "rng": rng}
      for shape in [(1, 1), (3, 4), (2, 10, 5), (2, 200, 100)]
      for dtype in float_types()
      for full_matrices in [False, True]
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("cpu")
  def testQr(self, shape, dtype, full_matrices, rng):
    m, n = shape[-2:]

    if full_matrices:
      mode, k = "complete", m
    else:
      mode, k = "reduced", min(m, n)

    a = rng(shape, dtype)
    lq, lr = np.linalg.qr(a, mode=mode)

    # onp.linalg.qr doesn't support broadcasting. But it seems like an
    # inevitable extension so we support it in our version.
    nq = onp.zeros(shape[:-2] + (m, k), dtype)
    nr = onp.zeros(shape[:-2] + (k, n), dtype)
    for index in onp.ndindex(*shape[:-2]):
      nq[index], nr[index] = onp.linalg.qr(a[index], mode=mode)

    max_rank = max(m, n)

    # Norm, adjusted for dimension and type.
    def norm(x):
      n = onp.linalg.norm(x, axis=(-2, -1))
      return n / (max_rank * onp.finfo(dtype).eps)

    def compare_orthogonal(q1, q2):
      # Q is unique up to sign, so normalize the sign first.
      sum_of_ratios = onp.sum(onp.divide(q1, q2), axis=-2, keepdims=True)
      phases = onp.divide(sum_of_ratios, onp.abs(sum_of_ratios))
      q1 *= phases
      self.assertTrue(onp.all(norm(q1 - q2) < 30))

    # Check a ~= qr
    self.assertTrue(onp.all(norm(a - onp.matmul(lq, lr)) < 30))

    # Compare the first 'k' vectors of Q; the remainder form an arbitrary
    # orthonormal basis for the null space.
    compare_orthogonal(nq[..., :k], lq[..., :k])

    # Check that q is close to unitary.
    self.assertTrue(onp.all(norm(onp.eye(k) - onp.matmul(T(lq), lq)) < 5))

    if not full_matrices and m >= n:
        jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype)),
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lhs_shape, rhs_shape in [
          ((1, 1), (1, 1)),
          ((4, 4), (4,)),
          ((8, 8), (8, 4)),
      ]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testSolve(self, lhs_shape, rhs_shape, dtype, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]

    self._CheckAgainstNumpy(onp.linalg.solve, np.linalg.solve, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.solve, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 4), (2, 5, 5), (200, 200), (5, 5, 5)]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testInv(self, shape, dtype, rng):
    def args_maker():
      invertible = False
      while not invertible:
        a = rng(shape, dtype)
        try:
          onp.linalg.inv(a)
          invertible = True
        except onp.linalg.LinAlgError:
          pass
      return [a]

    self._CheckAgainstNumpy(onp.linalg.inv, np.linalg.inv, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(np.linalg.inv, args_maker, check_dtypes=True)


class ScipyLinalgTest(jtu.JaxTestCase):

  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
       "shape": shape, "dtype": dtype, "rng": rng}
      for shape in [(1, 1), (4, 5), (10, 5), (50, 50)]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testLu(self, shape, dtype, rng):
    # TODO(phawkins): remove this after a jaxlib release.
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    args_maker = lambda: [rng(shape, dtype)]

    self._CheckAgainstNumpy(jsp.linalg.lu, osp.linalg.lu, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(jsp.linalg.lu, args_maker, check_dtypes=True)

  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
       "n": n, "dtype": dtype, "rng": rng}
      for n in [1, 4, 5, 200]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  @jtu.skip_on_devices("gpu", "tpu")
  def testLuFactor(self, n, dtype, rng):
    # TODO(phawkins): remove this after a jaxlib release.
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    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)
    self._CompileAndCheck(jsp.linalg.lu_factor, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_sym_pos={}_lower={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           sym_pos, lower),
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "sym_pos": sym_pos, "lower": lower, "rng": rng}
      for lhs_shape, rhs_shape in [
          ((1, 1), (1, 1)),
          ((4, 4), (4,)),
          ((8, 8), (8, 4)),
      ]
      for sym_pos, lower in [
        (False, False),
        (True, False),
        (True, True),
      ]
      for dtype in float_types() | complex_types()
      for rng in [jtu.rand_default()]))
  # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testSolve(self, lhs_shape, rhs_shape, dtype, sym_pos, lower, rng):
    if not hasattr(lapack, "jax_getrf"):
      self.skipTest("No LU implementation available")
    osp_fun = lambda lhs, rhs: osp.linalg.solve(lhs, rhs, sym_pos=sym_pos, lower=lower)
    jsp_fun = lambda lhs, rhs: jsp.linalg.solve(lhs, rhs, sym_pos=sym_pos, lower=lower)

    def args_maker():
      a = rng(lhs_shape, dtype)
      if sym_pos:
        a = onp.matmul(a, onp.conj(T(a)))
        a = onp.tril(a) if lower else onp.triu(a)
      return [a, rng(rhs_shape, dtype)]

    self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker,
                            check_dtypes=True, tol=1e-3)
    self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)


  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_lower={}_transposea={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           lower, transpose_a),
       "lower": lower, "transpose_a": transpose_a,
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lower, transpose_a in itertools.product([False, True], repeat=2)
      for lhs_shape, rhs_shape in [
          ((4, 4), (4,)),
          ((4, 4), (4, 3)),
          ((2, 8, 8), (2, 8, 10)),
      ]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testSolveTriangular(self, lower, transpose_a, lhs_shape,
                                 rhs_shape, dtype, rng):
    k = rng(lhs_shape, dtype)
    l = onp.linalg.cholesky(onp.matmul(k, T(k))
                            + lhs_shape[-1] * onp.eye(lhs_shape[-1]))
    l = l.astype(k.dtype)
    b = rng(rhs_shape, dtype)

    a = l if lower else T(l)
    inv = onp.linalg.inv(T(a) if transpose_a else a).astype(a.dtype)
    if len(lhs_shape) == len(rhs_shape):
      onp_ans = onp.matmul(inv, b)
    else:
      onp_ans = onp.einsum("...ij,...j->...i", inv, b)

    # The standard scipy.linalg.solve_triangular doesn't support broadcasting.
    # But it seems like an inevitable extension so we support it.
    ans = jsp.linalg.solve_triangular(
        l if lower else T(l), b, trans=1 if transpose_a else 0, lower=lower)

    self.assertAllClose(onp_ans, ans, check_dtypes=True)

  @parameterized.named_parameters(jtu.cases_from_list(
      {"testcase_name":
       "_lhs={}_rhs={}_lower={}_transposea={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),
           jtu.format_shape_dtype_string(rhs_shape, dtype),
           lower, transpose_a),
       "lower": lower, "transpose_a": transpose_a,
       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
       "rng": rng}
      for lower, transpose_a in itertools.product([False, True], repeat=2)
      for lhs_shape, rhs_shape in [
          ((4, 4), (4,)),
          ((4, 4), (4, 3)),
          ((2, 8, 8), (2, 8, 10)),
      ]
      for dtype in float_types()
      for rng in [jtu.rand_default()]))
  def testSolveTriangularGrad(self, lower, transpose_a, lhs_shape,
                                     rhs_shape, dtype, rng):
    # TODO(frostig): change ensemble to support a bigger rtol
    self.skipTest("rtol does not cover all devices and precision modes")
    A = np.tril(rng(lhs_shape, dtype) + 5 * onp.eye(lhs_shape[-1], dtype=dtype))
    A = A if lower else T(A)
    B = rng(rhs_shape, dtype)
    f = partial(jsp.linalg.solve_triangular, lower=lower,
                trans=1 if transpose_a else 0)
    jtu.check_grads(f, (A, B), 2, rtol=1e-3)

if __name__ == "__main__":
  absltest.main()
