@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 [[(2,), (2,), ((- 1), (- 1), (- 1), None)], [(2, 4), (2, 4), ((- 1), (- 1), (- 1), 0)], [(3, 4), (3, 4), ((- 1), (- 1), (- 1), 0)], [(3, 4), (3, 6, 5, 4), ((- 1), (- 1), (- 1), 0)], [(4, 3), (3, 6, 5, 4), (1, 0, (- 1), None)], [(6, 1, 3), (5, 3), ((- 1), (- 1), (- 1), None)], [(6, 1, 2), (5, 3), ((- 1), (- 1), (- 1), None)], [(10, 5, 2, 8), (1, 5, 1, 3), ((- 2), (- 1), (- 3), None)], [(4, 5, 2), (4, 5, 2), ((- 1), (- 1), 0, None)], [(4, 5, 2), (4, 5, 2), ((- 1), (- 1), (- 1), None)]] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testCross(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)])
    (axisa, axisb, axisc, axis) = axes
    jnp_fun = (lambda a, b: jnp.cross(a, b, axisa, axisb, axisc, axis))

    def onp_fun(a, b):
        a = (a.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else a)
        b = (b.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else b)
        out = onp.cross(a, b, axisa, axisb, axisc, axis)
        return out.astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    tol_spec = {dtypes.bfloat16: 0.3, onp.float16: 0.15}
    tol = max(jtu.tolerance(lhs_dtype, tol_spec), jtu.tolerance(rhs_dtype, tol_spec))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, 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), 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 [[(2,), (2,), ((- 1), (- 1), (- 1), None)], [(2, 4), (2, 4), ((- 1), (- 1), (- 1), 0)], [(3, 4), (3, 4), ((- 1), (- 1), (- 1), 0)], [(3, 4), (3, 6, 5, 4), ((- 1), (- 1), (- 1), 0)], [(4, 3), (3, 6, 5, 4), (1, 0, (- 1), None)], [(6, 1, 3), (5, 3), ((- 1), (- 1), (- 1), None)], [(6, 1, 2), (5, 3), ((- 1), (- 1), (- 1), None)], [(10, 5, 2, 8), (1, 5, 1, 3), ((- 2), (- 1), (- 3), None)], [(4, 5, 2), (4, 5, 2), ((- 1), (- 1), 0, None)], [(4, 5, 2), (4, 5, 2), ((- 1), (- 1), (- 1), None)]] for (lhs_dtype, rhs_dtype) in CombosWithReplacement(number_dtypes, 2))))
def testCross(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)])
    (axisa, axisb, axisc, axis) = axes
    jnp_fun = (lambda a, b: jnp.cross(a, b, axisa, axisb, axisc, axis))

    def onp_fun(a, b):
        a = (a.astype(onp.float32) if (lhs_dtype == jnp.bfloat16) else a)
        b = (b.astype(onp.float32) if (rhs_dtype == jnp.bfloat16) else b)
        out = onp.cross(a, b, axisa, axisb, axisc, axis)
        return out.astype(jnp.promote_types(lhs_dtype, rhs_dtype))
    tol_spec = {dtypes.bfloat16: 0.3, onp.float16: 0.15}
    tol = max(jtu.tolerance(lhs_dtype, tol_spec), jtu.tolerance(rhs_dtype, tol_spec))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, atol=tol, rtol=tol)