@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_inshape={}_axes=({},{})'.format(jtu.format_shape_dtype_string(arg_shape, dtype), ax1, ax2), 'arg_shape': arg_shape, 'dtype': dtype, 'ax1': ax1, 'ax2': ax2, 'rng_factory': jtu.rand_default} for (arg_shape, ax1, ax2) in [((3, 4), 0, 1), ((3, 4), 1, 0), ((3, 4, 5), 1, 2), ((3, 4, 5), (- 1), (- 2)), ((3, 4, 5), 0, 1)] for dtype in default_dtypes)))
def testSwapAxesStaticAxes(self, arg_shape, dtype, ax1, ax2, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda x: onp.swapaxes(x, ax1, ax2))
    jnp_fun = (lambda x: jnp.swapaxes(x, ax1, ax2))
    args_maker = (lambda : [rng(arg_shape, dtype)])
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_inshape={}_axes=({},{})'.format(jtu.format_shape_dtype_string(arg_shape, dtype), ax1, ax2), 'arg_shape': arg_shape, 'dtype': dtype, 'ax1': ax1, 'ax2': ax2, 'rng_factory': jtu.rand_default} for (arg_shape, ax1, ax2) in [((3, 4), 0, 1), ((3, 4), 1, 0), ((3, 4, 5), 1, 2), ((3, 4, 5), (- 1), (- 2)), ((3, 4, 5), 0, 1)] for dtype in default_dtypes)))
def testSwapAxesStaticAxes(self, arg_shape, dtype, ax1, ax2, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda x: onp.swapaxes(x, ax1, ax2))
    jnp_fun = (lambda x: jnp.swapaxes(x, ax1, ax2))
    args_maker = (lambda : [rng(arg_shape, dtype)])
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)