@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_{}_{}'.format(name, jtu.format_shape_dtype_string(lhs_shape, lhs_dtype), jtu.format_shape_dtype_string(rhs_shape, rhs_dtype)), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'rng_factory': rng_factory} for rng_factory in [jtu.rand_default] for (name, lhs_shape, rhs_shape) in [('vector-vector', (3,), (3,)), ('matrix-vector', (3, 3), (3,)), ('vector-matrix', (3,), (3, 3)), ('matrix-matrix', (3, 3), (3, 3)), ('vector-tensor', (3,), (5, 3, 2)), ('tensor-vector', (5, 3, 2), (2,)), ('matrix-tensor', (5, 2), (3, 2, 4)), ('tensor-matrix', (5, 2, 3), (3, 2)), ('tensor-tensor', (5, 3, 4), (5, 4, 1)), ('tensor-tensor-broadcast', (3, 1, 3, 4), (5, 4, 1))] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testMatmul(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng_factory):
    rng = rng_factory()

    def onp_fun(x, y):
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.matmul(x, y).astype(dtype)
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    tol = {onp.float16: 0.01, onp.float32: 0.02, onp.float64: 1e-12, onp.complex128: 1e-12}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.04
    self._CheckAgainstNumpy(onp_fun, jnp.matmul, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp.matmul, args_maker, check_dtypes=True, atol=tol, rtol=tol)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_{}_{}'.format(name, jtu.format_shape_dtype_string(lhs_shape, lhs_dtype), jtu.format_shape_dtype_string(rhs_shape, rhs_dtype)), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'rng_factory': rng_factory} for rng_factory in [jtu.rand_default] for (name, lhs_shape, rhs_shape) in [('vector-vector', (3,), (3,)), ('matrix-vector', (3, 3), (3,)), ('vector-matrix', (3,), (3, 3)), ('matrix-matrix', (3, 3), (3, 3)), ('vector-tensor', (3,), (5, 3, 2)), ('tensor-vector', (5, 3, 2), (2,)), ('matrix-tensor', (5, 2), (3, 2, 4)), ('tensor-matrix', (5, 2, 3), (3, 2)), ('tensor-tensor', (5, 3, 4), (5, 4, 1)), ('tensor-tensor-broadcast', (3, 1, 3, 4), (5, 4, 1))] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testMatmul(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng_factory):
    rng = rng_factory()

    def onp_fun(x, y):
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.matmul(x, y).astype(dtype)
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    tol = {onp.float16: 0.01, onp.float32: 0.02, onp.float64: 1e-12, onp.complex128: 1e-12}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.04
    self._CheckAgainstNumpy(onp_fun, jnp.matmul, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp.matmul, args_maker, check_dtypes=True, atol=tol, rtol=tol)