Closed copybara-service[bot] closed 2 years ago
[JAX] Update users of jax.ops.index... functions, which are deprecated.
jax.ops.index...
jax.ops.index[...]
jax.numpy.index_exp[...]
.at[...]
jax.ops.index_update(x, jax.ops.index[idx], y)
x.at[idx].set(y)
[JAX] Update users of
jax.ops.index...
functions, which are deprecated.jax.ops.index[...]
withjax.numpy.index_exp[...]
, which is a standard NumPy function that does the same thing.jax.ops.index[...]
, where the expression is passed directly to an indexed accessor function like.at[...]
.jax.ops.index_update(x, jax.ops.index[idx], y)
to use thex.at[idx].set(y)
APIs.