Closed ptigwe closed 2 years ago
Replacing the deprecated jax.ops.index_update with the suggested alternative of arr.at[idx].set(val). Another alternative which is to use masking tricks also yields the same effect is as follows:
jax.ops.index_update
arr.at[idx].set(val)
idx_axis0 = jnp.arange(prob_tensor.shape[0]) jnp.sum(idx_axis0 == jnp.expand_dims(idx_axis1, -1), 1)
LGTM! Thanks a lot @ptigwe!
cgarciae
commented
2 years ago
Replacing the deprecated
jax.ops.index_updatewith the suggested alternative ofarr.at[idx].set(val). Another alternative which is to use masking tricks also yields the same effect is as follows: