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","."]
code: def _update_diagonal(self, g, w, m, v1, v2, opt_params):
      learning_rate = opt_params["learning_rate"]
      beta2 = opt_params["second_moment_averaging"]
      weight_decay = opt_params["weight_decay"]
  
      is_beta2_1 = (beta2 == 1).astype(g.dtype)
      one_minus_beta2_except1 = is_beta2_1 + (1.0 - beta2) * (
          1.0 - is_beta2_1
      )
      v1[0] = beta2 * v1[0] + one_minus_beta2_except1 * g * g
  
      preconditioner = jnp.where(
          v1[0] > 0,
          1.0 / (jnp.sqrt(v1[0]) + 1e-16),
          jnp.zeros_like(v1[0]),
      )
  
      pg = preconditioner * g
      if self._graft:
          v2[0] += g * g
          preconditioner_graft = jnp.where(
              v2[0] > 0,
              1.0 / (jnp.sqrt(v2[0]) + 1e-16),
              jnp.zeros_like(v2[0]),
          )
          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, opt_params["momentum"]
          )
      else:
          update = pg
  
      w = w - (update * learning_rate).astype(w.dtype)
      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 4 for calling the SM3 obj's underlying member method code with example usage: obj._update_diagonal(g, w, m, v1, v2, opt_params) and returns: w,  m,  v1,  v2