@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_{}_{}'.format(jtu.format_shape_dtype_string(lhs_shape, lhs_dtype), jtu.format_shape_dtype_string(rhs_shape, rhs_dtype), axes), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'axes': axes, 'rng_factory': rng_factory} for rng_factory in [jtu.rand_default] for (lhs_shape, rhs_shape, axes) in [[(3,), (), 0], [(2, 3, 4), (5, 6, 7), 0], [(2, 3, 4), (3, 4, 5, 6), 2], [(2, 3, 4), (5, 4, 3, 6), [1, 2]], [(2, 3, 4), (5, 4, 3, 6), [[1, 2], [2, 1]]], [(1, 2, 3, 4), (4, 5, 3, 6), [[2, 3], [2, 0]]]] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testTensordot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, axes, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    jnp_fun = (lambda a, b: jnp.tensordot(a, b, axes))

    def onp_fun(a, b):
        a = (a if (lhs_dtype != jnp.bfloat16) else a.astype(onp.float32))
        b = (b if (rhs_dtype != jnp.bfloat16) else b.astype(onp.float32))
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.tensordot(a, b, axes).astype(dtype)
    tol = {onp.float16: 0.1, onp.float32: 0.001, onp.float64: 1e-12, onp.complex64: 0.001, onp.complex128: 1e-12}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.2
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_{}_{}'.format(jtu.format_shape_dtype_string(lhs_shape, lhs_dtype), jtu.format_shape_dtype_string(rhs_shape, rhs_dtype), axes), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'axes': axes, 'rng_factory': rng_factory} for rng_factory in [jtu.rand_default] for (lhs_shape, rhs_shape, axes) in [[(3,), (), 0], [(2, 3, 4), (5, 6, 7), 0], [(2, 3, 4), (3, 4, 5, 6), 2], [(2, 3, 4), (5, 4, 3, 6), [1, 2]], [(2, 3, 4), (5, 4, 3, 6), [[1, 2], [2, 1]]], [(1, 2, 3, 4), (4, 5, 3, 6), [[2, 3], [2, 0]]]] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testTensordot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, axes, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)])
    jnp_fun = (lambda a, b: jnp.tensordot(a, b, axes))

    def onp_fun(a, b):
        a = (a if (lhs_dtype != jnp.bfloat16) else a.astype(onp.float32))
        b = (b if (rhs_dtype != jnp.bfloat16) else b.astype(onp.float32))
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.tensordot(a, b, axes).astype(dtype)
    tol = {onp.float16: 0.1, onp.float32: 0.001, onp.float64: 1e-12, onp.complex64: 0.001, onp.complex128: 1e-12}
    if (jtu.device_under_test() == 'tpu'):
        tol[onp.float32] = tol[onp.complex64] = 0.2
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)