@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 arg_dtypes in CombosWithReplacement(default_dtypes, 2) for base_shape in [(4,), (3, 4), (2, 3, 4)] for axis in range(((- len(base_shape)) + 1), len(base_shape)))))
def testAppend(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(arr, values):
        arr = (arr.astype(onp.float32) if (arr.dtype == jnp.bfloat16) else arr)
        values = (values.astype(onp.float32) if (values.dtype == jnp.bfloat16) else values)
        out = onp.append(arr, values, axis=axis)
        return out.astype(jnp.promote_types(*arg_dtypes))
    jnp_fun = (lambda arr, values: jnp.append(arr, values, 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 arg_dtypes in CombosWithReplacement(default_dtypes, 2) for base_shape in [(4,), (3, 4), (2, 3, 4)] for axis in range(((- len(base_shape)) + 1), len(base_shape)))))
def testAppend(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(arr, values):
        arr = (arr.astype(onp.float32) if (arr.dtype == jnp.bfloat16) else arr)
        values = (values.astype(onp.float32) if (values.dtype == jnp.bfloat16) else values)
        out = onp.append(arr, values, axis=axis)
        return out.astype(jnp.promote_types(*arg_dtypes))
    jnp_fun = (lambda arr, values: jnp.append(arr, values, 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)