@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_axis={}_{}sections'.format(jtu.format_shape_dtype_string(shape, dtype), axis, num_sections), 'shape': shape, 'num_sections': num_sections, 'axis': axis, 'dtype': dtype, 'rng_factory': jtu.rand_default} for (shape, axis, num_sections) in [((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((12, 4), 1, 2), ((2, 3, 4), (- 1), 2), ((2, 3, 4), (- 2), 3)] for dtype in default_dtypes)))
def testSplitStaticInt(self, shape, num_sections, axis, dtype, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda x: onp.split(x, num_sections, axis=axis))
    jnp_fun = (lambda x: jnp.split(x, num_sections, axis=axis))
    args_maker = (lambda : [rng(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': '_{}_axis={}_{}sections'.format(jtu.format_shape_dtype_string(shape, dtype), axis, num_sections), 'shape': shape, 'num_sections': num_sections, 'axis': axis, 'dtype': dtype, 'rng_factory': jtu.rand_default} for (shape, axis, num_sections) in [((3,), 0, 3), ((12,), 0, 3), ((12, 4), 0, 4), ((12, 4), 1, 2), ((2, 3, 4), (- 1), 2), ((2, 3, 4), (- 2), 3)] for dtype in default_dtypes)))
def testSplitStaticInt(self, shape, num_sections, axis, dtype, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda x: onp.split(x, num_sections, axis=axis))
    jnp_fun = (lambda x: jnp.split(x, num_sections, axis=axis))
    args_maker = (lambda : [rng(shape, dtype)])
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)