Taking a gradient of a parameter input into index_update

[sunilkpai]

In jax, it appears that by default, when a tensor is assigned to traced parameters, it is copied into the original tensor and no longer traced. Is this the right interpretation?

Code to reproduce:

import jax.numpy as jnp
from jax import value_and_grad, jit

y = jnp.linspace(0, 10, 11)

def idx_update(x):
    m = y
    y.at[2:8].set(x)
    return y[:5].sum()

x = jnp.linspace(0, 5, 6)
grad_fn = jit(value_and_grad(idx_update))
grad_fn(x)

Output:

(DeviceArray(10., dtype=float64),
 DeviceArray([0., 0., 0., 0., 0., 0.], dtype=float64))

My desired output:

(DeviceArray(4., dtype=float64),
 DeviceArray([1., 1., 1., 0., 0., 0.], dtype=float64))

Is there some way I can achieve the desired output (assuming that's not JAX's intent / this is not a bug)?

[davisyoshida]

The line y.at[2:8].set(x) currently has no effect. Presumably you want:

z = y.at[2:8].set(x)
return z[:5].sum()
Shouldn't the desired sum be 4, not 10 though? Unless I'm misreading the intent.

[sunilkpai]

Oops, you are correct on the desired sum being 4 (edited this in the original question). Thanks for your help!