@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 [('matrix-scalar', (3, 3), ()), ('scalar-matrix', (), (3, 3)), ('matrix-vector', (4, 5), (5,)), ('vector-matrix', (6,), (6, 4)), ('matrix-matrix', (3, 4), (4, 5)), ('tensor-vector', (4, 3, 2), (2,)), ('vector-tensor', (2,), (3, 2, 4)), ('tensor-matrix', (4, 3, 2), (2, 5)), ('matrix-tensor', (5, 2), (3, 2, 4)), ('tensor-tensor', (2, 3, 4), (5, 4, 1))] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testDot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    tol = {onp.float16: 0.01, onp.float32: 1e-05, onp.float64: 1e-14, onp.complex128: 1e-14}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.2

    def onp_dot(x, y):
        x = (x.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else x)
        y = (y.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else y)
        return onp.dot(x, y).astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    self._CheckAgainstNumpy(onp_dot, jnp.dot, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp.dot, 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 [('matrix-scalar', (3, 3), ()), ('scalar-matrix', (), (3, 3)), ('matrix-vector', (4, 5), (5,)), ('vector-matrix', (6,), (6, 4)), ('matrix-matrix', (3, 4), (4, 5)), ('tensor-vector', (4, 3, 2), (2,)), ('vector-tensor', (2,), (3, 2, 4)), ('tensor-matrix', (4, 3, 2), (2, 5)), ('matrix-tensor', (5, 2), (3, 2, 4)), ('tensor-tensor', (2, 3, 4), (5, 4, 1))] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testDot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    tol = {onp.float16: 0.01, onp.float32: 1e-05, onp.float64: 1e-14, onp.complex128: 1e-14}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.2

    def onp_dot(x, y):
        x = (x.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else x)
        y = (y.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else y)
        return onp.dot(x, y).astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    self._CheckAgainstNumpy(onp_dot, jnp.dot, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp.dot, args_maker, check_dtypes=True, atol=tol, rtol=tol)