@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_k={}_axes={}'.format(jtu.format_shape_dtype_string(shape, dtype), k, axes), 'rng_factory': rng_factory, 'shape': shape, 'dtype': dtype, 'k': k, 'axes': axes} for (shape, axes) in [[(2, 3), (0, 1)], [(2, 3), (1, 0)], [(4, 3, 2), (0, 2)], [(4, 3, 2), (2, 1)]] for k in range((- 3), 4) for dtype in default_dtypes for rng_factory in [jtu.rand_default])))
def testRot90(self, shape, dtype, k, axes, rng_factory):
    rng = rng_factory()
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = (lambda x: jnp.rot90(x, k, axes))
    onp_op = (lambda x: onp.rot90(x, k, axes))
    self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_k={}_axes={}'.format(jtu.format_shape_dtype_string(shape, dtype), k, axes), 'rng_factory': rng_factory, 'shape': shape, 'dtype': dtype, 'k': k, 'axes': axes} for (shape, axes) in [[(2, 3), (0, 1)], [(2, 3), (1, 0)], [(4, 3, 2), (0, 2)], [(4, 3, 2), (2, 1)]] for k in range((- 3), 4) for dtype in default_dtypes for rng_factory in [jtu.rand_default])))
def testRot90(self, shape, dtype, k, axes, rng_factory):
    rng = rng_factory()
    args_maker = self._GetArgsMaker(rng, [shape], [dtype])
    jnp_op = (lambda x: jnp.rot90(x, k, axes))
    onp_op = (lambda x: onp.rot90(x, k, axes))
    self._CheckAgainstNumpy(onp_op, jnp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)