In place matrix/tensor update

[TuanNguyen27]

I know that JAX explicitly doesn't support in place operation, but i saw this argument about memory location in Dive Into Deep Learning book. I'm wondering if there's a way to achieve this in JAX, e.g update a tensor and keep the same memory location.

The standard JAX way to do the above code would be (i think)

Z = jnp.zeros_like(Y)
print('id(Z):', id(Z))
Z = jax.ops.index_update(Z, jax.ops.index[:], X+Y)
print('id(Z):', id(Z))

but as you can tell, it does not retain the original memory location

id(Z): 140308028317280
id(Z): 140308028316752

Would love to hear more about this! Is the issue about memory reference being stale valid, and how do JAX frameworks tackle this ?

[jakevdp]

See #1733 for a version of this request and some recent updates.