@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_axis={}_baseshape=[{}]_dtypes=[{}]'.format(axis, ','.join((str(d) for d in base_shape)), ','.join((onp.dtype(dtype).name for dtype in arg_dtypes))), 'axis': axis, 'base_shape': base_shape, 'arg_dtypes': arg_dtypes, 'rng_factory': jtu.rand_default} for num_arrs in [3] for arg_dtypes in CombosWithReplacement(default_dtypes, num_arrs) for base_shape in [(4,), (3, 4), (2, 3, 4)] for axis in range(((- len(base_shape)) + 1), len(base_shape)))))
def testConcatenate(self, axis, base_shape, arg_dtypes, rng_factory):
    rng = rng_factory()
    wrapped_axis = (axis % len(base_shape))
    shapes = [((base_shape[:wrapped_axis] + (size,)) + base_shape[(wrapped_axis + 1):]) for (size, _) in zip(itertools.cycle([3, 1, 4]), arg_dtypes)]

    def onp_fun(*args):
        args = [(x if (x.dtype != jnp.bfloat16) else x.astype(onp.float32)) for x in args]
        dtype = functools.reduce(jnp.promote_types, arg_dtypes)
        return onp.concatenate(args, axis=axis).astype(dtype)
    jnp_fun = (lambda *args: jnp.concatenate(args, axis=axis))

    def args_maker():
        return [rng(shape, dtype) for (shape, dtype) in zip(shapes, arg_dtypes)]
    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={}_baseshape=[{}]_dtypes=[{}]'.format(axis, ','.join((str(d) for d in base_shape)), ','.join((onp.dtype(dtype).name for dtype in arg_dtypes))), 'axis': axis, 'base_shape': base_shape, 'arg_dtypes': arg_dtypes, 'rng_factory': jtu.rand_default} for num_arrs in [3] for arg_dtypes in CombosWithReplacement(default_dtypes, num_arrs) for base_shape in [(4,), (3, 4), (2, 3, 4)] for axis in range(((- len(base_shape)) + 1), len(base_shape)))))
def testConcatenate(self, axis, base_shape, arg_dtypes, rng_factory):
    rng = rng_factory()
    wrapped_axis = (axis % len(base_shape))
    shapes = [((base_shape[:wrapped_axis] + (size,)) + base_shape[(wrapped_axis + 1):]) for (size, _) in zip(itertools.cycle([3, 1, 4]), arg_dtypes)]

    def onp_fun(*args):
        args = [(x if (x.dtype != jnp.bfloat16) else x.astype(onp.float32)) for x in args]
        dtype = functools.reduce(jnp.promote_types, arg_dtypes)
        return onp.concatenate(args, axis=axis).astype(dtype)
    jnp_fun = (lambda *args: jnp.concatenate(args, axis=axis))

    def args_maker():
        return [rng(shape, dtype) for (shape, dtype) in zip(shapes, arg_dtypes)]
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)