@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)), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'rng_factory': jtu.rand_default} for (lhs_shape, lhs_dtype) in _shape_and_dtypes(all_shapes, inexact_dtypes) for (rhs_shape, rhs_dtype) in _shape_and_dtypes(all_shapes, inexact_dtypes) if ((len(jtu._dims_of_shape(lhs_shape)) == 0) or (len(jtu._dims_of_shape(rhs_shape)) == 0) or (lhs_shape[(- 1)] == rhs_shape[(- 1)])))))
def testInner(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)])

    def onp_fun(lhs, rhs):
        lhs = (lhs if (lhs_dtype != jnp.bfloat16) else lhs.astype(onp.float32))
        rhs = (rhs if (rhs_dtype != jnp.bfloat16) else rhs.astype(onp.float32))
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.inner(lhs, rhs).astype(dtype)
    jnp_fun = (lambda lhs, rhs: jnp.inner(lhs, rhs))
    tol_spec = {onp.float16: 0.01, onp.float32: 1e-05, onp.float64: 1e-13}
    if (jtu.device_under_test() == 'tpu'):
        tol_spec[onp.float32] = tol_spec[onp.complex64] = 0.2
    tol = max(jtu.tolerance(lhs_dtype, tol_spec), jtu.tolerance(rhs_dtype, tol_spec))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=False, atol=tol, rtol=tol)
@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)), 'lhs_shape': lhs_shape, 'lhs_dtype': lhs_dtype, 'rhs_shape': rhs_shape, 'rhs_dtype': rhs_dtype, 'rng_factory': jtu.rand_default} for (lhs_shape, lhs_dtype) in _shape_and_dtypes(all_shapes, inexact_dtypes) for (rhs_shape, rhs_dtype) in _shape_and_dtypes(all_shapes, inexact_dtypes) if ((len(jtu._dims_of_shape(lhs_shape)) == 0) or (len(jtu._dims_of_shape(rhs_shape)) == 0) or (lhs_shape[(- 1)] == rhs_shape[(- 1)])))))
def testInner(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)])

    def onp_fun(lhs, rhs):
        lhs = (lhs if (lhs_dtype != jnp.bfloat16) else lhs.astype(onp.float32))
        rhs = (rhs if (rhs_dtype != jnp.bfloat16) else rhs.astype(onp.float32))
        dtype = jnp.promote_types(lhs_dtype, rhs_dtype)
        return onp.inner(lhs, rhs).astype(dtype)
    jnp_fun = (lambda lhs, rhs: jnp.inner(lhs, rhs))
    tol_spec = {onp.float16: 0.01, onp.float32: 1e-05, onp.float64: 1e-13}
    if (jtu.device_under_test() == 'tpu'):
        tol_spec[onp.float32] = tol_spec[onp.complex64] = 0.2
    tol = max(jtu.tolerance(lhs_dtype, tol_spec), jtu.tolerance(rhs_dtype, tol_spec))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=False, atol=tol, rtol=tol)