def testLu(self, shape, dtype):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    x, = args_maker()
    p, l, u = jsp.linalg.lu(x)
    self.assertAllClose(x, np.matmul(p, np.matmul(l, u)),
                        rtol={np.float32: 1e-3, np.float64: 1e-12,
                              np.complex64: 1e-3, np.complex128: 1e-12})
    self._CompileAndCheck(jsp.linalg.lu, args_maker)