[JAX] Replace uses of deprecated jax.ops.index_update(x, idx, y) APIs with their up-to-date, more succinct equivalent x.at[idx].set(y).

The JAX operators:
jax.ops.index_update(x, jax.ops.index[idx], y)
jax.ops.index_add(x, jax.ops.index[idx], y)
...

have long been deprecated in lieu of their more succinct counterparts:
x.at[idx].set(y)
x.at[idx].add(y)
...

This change updates users of the deprecated APIs to use the current APIs, in preparation for removing the deprecated forms from JAX.

The main subtlety is that if `x` is not a JAX array, we must cast it to one using `jnp.asarray(x)` before using the new form, since `.at[...]` is only defined on JAX arrays.

PiperOrigin-RevId: 400214219

Author: trajax authors

trajax/optimizers.py

@@ -344,15 +344,15 @@ def ddp_rollout(dynamics, X, U, K, k, alpha, *args):
   T, m = U.shape
   Xnew = np.zeros((T + 1, n))
   Unew = np.zeros((T, m))
-  Xnew = ops.index_update(Xnew, ops.index[0], X[0])
+  Xnew = Xnew.at[0].set(X[0])
 
   def body(t, inputs):
     Xnew, Unew = inputs
     del_u = alpha * k[t] + np.matmul(K[t], Xnew[t] - X[t])
     u = U[t] + del_u
     x = dynamics(Xnew[t], u, t, *args)
-    Unew = ops.index_update(Unew, ops.index[t], u)
-    Xnew = ops.index_update(Xnew, ops.index[t + 1], x)
+    Unew = Unew.at[t].set(u)
+    Xnew = Xnew.at[t + 1].set(x)
     return Xnew, Unew
 
   return lax.fori_loop(0, T, body, (Xnew, Unew))
@@ -742,9 +742,8 @@ def gaussian_samples(random_key, mean, stdev, control_low, control_high,
   smoothing_coef = hyperparams['sampling_smoothing']
 
   def body_fun(t, noises):
-    return jax.ops.index_update(
-        noises, jax.ops.index[:, t], smoothing_coef * noises[:, t - 1] +
-        np.sqrt(1 - smoothing_coef**2) * noises[:, t])
+    return noises.at[:, t].set(smoothing_coef * noises[:, t - 1] +
+                               np.sqrt(1 - smoothing_coef**2) * noises[:, t])
 
   noises = jax.lax.fori_loop(1, horizon, body_fun, noises)
   samples = noises * stdev

trajax/tvlqr.py

@@ -46,14 +46,14 @@ def rollout(K, k, x0, A, B, c):
   T, m, n = K.shape
   X = np.zeros((T + 1, n))
   U = np.zeros((T, m))
-  X = ops.index_update(X, ops.index[0], x0)
+  X = X.at[0].set(x0)
 
   def body(t, inputs):
     X, U = inputs
     u = np.matmul(K[t], X[t]) + k[t]
     x = np.matmul(A[t], X[t]) + np.matmul(B[t], u) + c[t]
-    X = ops.index_update(X, ops.index[t + 1], x)
-    U = ops.index_update(U, ops.index[t], u)
+    X = X.at[t + 1].set(x)
+    U = U.at[t].set(u)
     return X, U
 
   return lax.fori_loop(0, T, body, (X, U))
@@ -147,10 +147,10 @@ def body(tt, inputs):
     t = T - 1 - tt
     P_t, p_t, K_t, k_t = lqr_step(P[t+1], p[t+1], Q[t], q[t], R[t], r[t], M[t],
                                   A[t], B[t], c[t])
-    K = ops.index_update(K, ops.index[t], K_t)
-    k = ops.index_update(k, ops.index[t], k_t)
-    P = ops.index_update(P, ops.index[t], P_t)
-    p = ops.index_update(p, ops.index[t], p_t)
+    K = K.at[t].set(K_t)
+    k = k.at[t].set(k_t)
+    P = P.at[t].set(P_t)
+    p = p.at[t].set(p_t)
 
     return K, k, P, p