@parameterized.named_parameters(itertools.chain.from_iterable((jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix(rec.test_name, shapes, dtypes), 'rng_factory': rec.rng_factory, 'shapes': shapes, 'dtypes': dtypes, 'name': rec.name, 'tol': rec.tolerance} for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(rec.shapes, rec.nargs)) for dtypes in itertools.product(*(_valid_dtypes_for_shape(s, rec.dtypes) for s in shapes)))) for rec in JAX_OPERATOR_OVERLOADS)))
def testOperatorOverload(self, name, rng_factory, shapes, dtypes, tol):
    rng = rng_factory()
    args_maker = self._GetArgsMaker(rng, shapes, dtypes, onp_arrays=False)
    fun = (lambda *xs: getattr(operator, name.strip('_'))(*xs))
    scalar_arg = ((jtu.PYTHON_SCALAR_SHAPE in shapes) or (jtu.NUMPY_SCALAR_SHAPE in shapes) or (() in shapes))
    empty_shape = any(((isinstance(s, tuple) and (0 in s)) for s in shapes))
    self._CompileAndCheck(fun, args_maker, check_dtypes=True, atol=tol, rtol=tol)
@parameterized.named_parameters(itertools.chain.from_iterable((jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix(rec.test_name, shapes, dtypes), 'rng_factory': rec.rng_factory, 'shapes': shapes, 'dtypes': dtypes, 'name': rec.name, 'tol': rec.tolerance} for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(rec.shapes, rec.nargs)) for dtypes in itertools.product(*(_valid_dtypes_for_shape(s, rec.dtypes) for s in shapes)))) for rec in JAX_OPERATOR_OVERLOADS)))
def testOperatorOverload(self, name, rng_factory, shapes, dtypes, tol):
    rng = rng_factory()
    args_maker = self._GetArgsMaker(rng, shapes, dtypes, onp_arrays=False)
    fun = (lambda *xs: getattr(operator, name.strip('_'))(*xs))
    scalar_arg = ((jtu.PYTHON_SCALAR_SHAPE in shapes) or (jtu.NUMPY_SCALAR_SHAPE in shapes) or (() in shapes))
    empty_shape = any(((isinstance(s, tuple) and (0 in s)) for s in shapes))
    self._CompileAndCheck(fun, args_maker, check_dtypes=True, atol=tol, rtol=tol)