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