@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_ishape={}_axis={}'.format(jtu.format_shape_dtype_string(x_shape, dtype), i_shape, axis), 'rng_factory': rng_factory, 'x_shape': x_shape, 'i_shape': i_shape, 'dtype': dtype, 'axis': axis} for (x_shape, i_shape) in filter(_shapes_are_equal_length, filter(_shapes_are_broadcast_compatible, CombosWithReplacement(nonempty_nonscalar_array_shapes, 2))) for axis in itertools.chain(range(len(x_shape)), [(- 1)], [cast(Optional[int], None)]) for dtype in default_dtypes for rng_factory in [jtu.rand_default])))
def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
    rng = rng_factory()
    i_shape = onp.array(i_shape)
    if (axis is None):
        i_shape = [onp.prod(i_shape, dtype=onp.int64)]
    else:
        i_shape[axis] *= 3
        i_shape = list(i_shape)

    def args_maker():
        x = rng(x_shape, dtype)
        n = (onp.prod(x_shape, dtype=onp.int32) if (axis is None) else x_shape[axis])
        i = ((rng(i_shape, onp.int32) % ((2 * n) - 1)) - (n - 1))
        return (x, i)
    jnp_op = (lambda x, i: jnp.take_along_axis(x, i, axis=axis))
    if hasattr(onp, 'take_along_axis'):
        onp_op = (lambda x, i: onp.take_along_axis(x, i, axis=axis))
        self._CheckAgainstNumpy(jnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_{}_ishape={}_axis={}'.format(jtu.format_shape_dtype_string(x_shape, dtype), i_shape, axis), 'rng_factory': rng_factory, 'x_shape': x_shape, 'i_shape': i_shape, 'dtype': dtype, 'axis': axis} for (x_shape, i_shape) in filter(_shapes_are_equal_length, filter(_shapes_are_broadcast_compatible, CombosWithReplacement(nonempty_nonscalar_array_shapes, 2))) for axis in itertools.chain(range(len(x_shape)), [(- 1)], [cast(Optional[int], None)]) for dtype in default_dtypes for rng_factory in [jtu.rand_default])))
def testTakeAlongAxis(self, x_shape, i_shape, dtype, axis, rng_factory):
    rng = rng_factory()
    i_shape = onp.array(i_shape)
    if (axis is None):
        i_shape = [onp.prod(i_shape, dtype=onp.int64)]
    else:
        i_shape[axis] *= 3
        i_shape = list(i_shape)

    def args_maker():
        x = rng(x_shape, dtype)
        n = (onp.prod(x_shape, dtype=onp.int32) if (axis is None) else x_shape[axis])
        i = ((rng(i_shape, onp.int32) % ((2 * n) - 1)) - (n - 1))
        return (x, i)
    jnp_op = (lambda x, i: jnp.take_along_axis(x, i, axis=axis))
    if hasattr(onp, 'take_along_axis'):
        onp_op = (lambda x, i: onp.take_along_axis(x, i, axis=axis))
        self._CheckAgainstNumpy(jnp_op, onp_op, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_op, args_maker, check_dtypes=True)