When using the jax_array.at[idx] function, I wish to be able to set values at both a set of specified rows and columns within the jax_array to another jax_array containing values in the same shape. For example, given a 5x5 jax array, I might want to set the values, jax_array.at[[0,3],:][:,[1,2]] to some 2x2 array of values. However, I am coming across an issue where the _IndexUpdateRef' object is not subscriptable. I understand the idea of the error (and I get a similar one when using 2 chained .at[]s), but I want to know if there is anyway to achieve the desired functionality within 1 line. Continue reading...
