def testTakeAlongAxisIssue1521(self):
    idx = jnp.repeat(jnp.arange(3), 10).reshape((30, 1))

    def f(x):
        y = (x * jnp.arange(3.0).reshape((1, 3)))
        return jnp.take_along_axis(y, idx, (- 1)).sum()
    check_grads(f, (1.0,), order=1)
def testTakeAlongAxisIssue1521(self):
    idx = jnp.repeat(jnp.arange(3), 10).reshape((30, 1))

    def f(x):
        y = (x * jnp.arange(3.0).reshape((1, 3)))
        return jnp.take_along_axis(y, idx, (- 1)).sum()
    check_grads(f, (1.0,), order=1)