@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_shifts={}_axis={}'.format(jtu.format_shape_dtype_string(shape, dtype), shifts, axis), 'rng_factory': rng_factory, 'shape': shape, 'dtype': dtype, 'shifts': shifts, 'axis': axis} for dtype in all_dtypes for shape in [(3, 4), (3, 4, 5), (7, 4, 0)] for (shifts, axis) in [(3, None), (1, 1), ((3,), (0,)), (((- 2),), ((- 2),)), ((1, 2), (0, (- 1))), ((4, 2, 5, 5, 2, 4), None), (100, None)] for rng_factory in [jtu.rand_default])))
def testRoll(self, shape, dtype, shifts, axis, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(shape, dtype), onp.array(shifts)])
    jnp_op = partial(jnp.roll, axis=axis)
    onp_op = partial(onp.roll, axis=axis)
    self._CheckAgainstNumpy(jnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_shifts={}_axis={}'.format(jtu.format_shape_dtype_string(shape, dtype), shifts, axis), 'rng_factory': rng_factory, 'shape': shape, 'dtype': dtype, 'shifts': shifts, 'axis': axis} for dtype in all_dtypes for shape in [(3, 4), (3, 4, 5), (7, 4, 0)] for (shifts, axis) in [(3, None), (1, 1), ((3,), (0,)), (((- 2),), ((- 2),)), ((1, 2), (0, (- 1))), ((4, 2, 5, 5, 2, 4), None), (100, None)] for rng_factory in [jtu.rand_default])))
def testRoll(self, shape, dtype, shifts, axis, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [rng(shape, dtype), onp.array(shifts)])
    jnp_op = partial(jnp.roll, axis=axis)
    onp_op = partial(onp.roll, axis=axis)
    self._CheckAgainstNumpy(jnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)