@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_axis={}'.format(jtu.format_test_name_suffix('', ([shape] * len(dtypes)), dtypes), axis), 'shape': shape, 'axis': axis, 'dtypes': dtypes, 'rng_factory': rng_factory} for dtypes in [[onp.float32], [onp.float32, onp.float32], [onp.float32, onp.int32, onp.float32], [onp.float32, onp.int64, onp.float32], [onp.float32, onp.int32, onp.float64]] for shape in [(), (2,), (3, 4), (1, 100)] for axis in range((- len(shape)), (len(shape) + 1)) for rng_factory in [jtu.rand_default])))
def testStack(self, shape, axis, dtypes, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [[rng(shape, dtype) for dtype in dtypes]])
    onp_fun = _promote_like_jnp(partial(onp.stack, axis=axis))
    jnp_fun = partial(jnp.stack, axis=axis)
    self._CheckAgainstNumpy(jnp_fun, onp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_axis={}'.format(jtu.format_test_name_suffix(, ([shape] * len(dtypes)), dtypes), axis), 'shape': shape, 'axis': axis, 'dtypes': dtypes, 'rng_factory': rng_factory} for dtypes in [[onp.float32], [onp.float32, onp.float32], [onp.float32, onp.int32, onp.float32], [onp.float32, onp.int64, onp.float32], [onp.float32, onp.int32, onp.float64]] for shape in [(), (2,), (3, 4), (1, 100)] for axis in range((- len(shape)), (len(shape) + 1)) for rng_factory in [jtu.rand_default])))
def testStack(self, shape, axis, dtypes, rng_factory):
    rng = rng_factory()
    args_maker = (lambda : [[rng(shape, dtype) for dtype in dtypes]])
    onp_fun = _promote_like_jnp(partial(onp.stack, axis=axis))
    jnp_fun = partial(jnp.stack, axis=axis)
    self._CheckAgainstNumpy(jnp_fun, onp_fun, args_maker, check_dtypes=True)