class_docstr: SM3 optimizer, as described in https://arxiv.org/abs/1901.11150.
class_docstr_tok: ["SM3","optimizer",",","as","described","in","https",":","\/\/arxiv.org\/abs\/1901.11150","."]
docstr: Update for higher-rank parameters.
docstr_tok: ["Update","for","higher-rank","parameters","."]
code: def _update_sketched(self, g, w, m, v1, v2, opt_params):
      """Update for higher-rank parameters."""
      learning_rate = opt_params["learning_rate"]
      momentum = opt_params["momentum"]
      beta2 = opt_params["second_moment_averaging"]
      weight_decay = opt_params["weight_decay"]
  
      shape = w.shape
      rank = len(shape)
      reshaped_accumulators = [
          jnp.reshape(v1[i], self._expanded_shape(shape, i))
          for i in range(rank)
      ]
      acc = self._minimum(reshaped_accumulators)
  
      is_beta2_1 = (beta2 == 1).astype(g.dtype)
      one_minus_beta2_except1 = is_beta2_1 + (1.0 - beta2) * (
          1.0 - is_beta2_1
      )
      acc = beta2 * acc + one_minus_beta2_except1 * g * g
  
      preconditioner = jnp.where(
          acc > 0.0, 1.0 / (jnp.sqrt(acc) + 1e-16), jnp.zeros_like(acc)
      )
      pg = g * preconditioner
      if self._graft:
          v2_acc = self._minimum(
              [
                  jnp.reshape(v2[i], self._expanded_shape(shape, i))
                  for i in range(rank)
              ]
          )
          v2_acc = v2_acc + g * g
          preconditioner_graft = jnp.where(
              v2_acc > 0.0,
              1.0 / (jnp.sqrt(v2_acc) + 1e-16),
              jnp.zeros_like(v2_acc),
          )
          pg_graft = preconditioner_graft * g
          pg_norm = jnp.linalg.norm(pg)
          pg_graft_norm = jnp.linalg.norm(pg_graft)
          pg = pg * (pg_graft_norm / (pg_norm + 1e-16))
  
      pg = pg + w * weight_decay
  
      if self._has_momentum:
          m, update = self._momentum_update(pg, m, momentum)
      else:
          update = pg
  
      w = w - (learning_rate * update).astype(w.dtype)
      for i in range(len(v1)):
          axes = list(range(int(i))) + list(range(int(i) + 1, rank))
          dim_accumulator = jnp.amax(acc, axis=axes)
          v1[i] = dim_accumulator
  
      if self._graft:
          for i in range(len(v2)):
              axes = list(range(int(i))) + list(range(int(i) + 1, rank))
              dim_accumulator = jnp.amax(v2_acc, axis=axes)
              v2[i] = dim_accumulator
      return w, (m, v1, v2)
  
file_path: trax/trax/optimizers/sm3.py
imports: import enum
  from trax.fastmath import numpy
  from trax.optimizers import base
label_desc: Use image node_id 7 for calling the SM3 obj's underlying member method code with example usage: obj._update_sketched(g, w, m, v1, v2, opt_params) and returns: w,  m,  v1,  v2