@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, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name)} for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(rec.shapes, rec.nargs)) for dtypes in filter(_dtypes_are_compatible_for_bitwise_ops, CombosWithReplacement(rec.dtypes, rec.nargs)))) for rec in JAX_BITWISE_OP_RECORDS)))
def testBitwiseOp(self, onp_op, jnp_op, rng_factory, shapes, dtypes):
    rng = rng_factory()
    if ((not FLAGS.jax_enable_x64) and any(((jnp.iinfo(dtype).bits == 64) for dtype in dtypes))):
        self.skipTest('x64 types are disabled by jax_enable_x64')
    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
    self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=(jtu.PYTHON_SCALAR_SHAPE not in shapes))
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
@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, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name)} for shapes in filter(_shapes_are_broadcast_compatible, CombosWithReplacement(rec.shapes, rec.nargs)) for dtypes in filter(_dtypes_are_compatible_for_bitwise_ops, CombosWithReplacement(rec.dtypes, rec.nargs)))) for rec in JAX_BITWISE_OP_RECORDS)))
def testBitwiseOp(self, onp_op, jnp_op, rng_factory, shapes, dtypes):
    rng = rng_factory()
    if ((not FLAGS.jax_enable_x64) and any(((jnp.iinfo(dtype).bits == 64) for dtype in dtypes))):
        self.skipTest('x64 types are disabled by jax_enable_x64')
    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
    self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=(jtu.PYTHON_SCALAR_SHAPE not in shapes))
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)