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 init(self, w):
      momentum = []
      if self._has_momentum:
          momentum = jnp.zeros_like(w)
      v1s = [jnp.zeros(sz, dtype=w.dtype) for sz in w.shape]
      v2s = []
      if self._graft:
          v2s = [jnp.zeros(sz, dtype=w.dtype) for sz in w.shape]
      return (momentum, v1s, v2s)
  
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 2 for calling the SM3 obj's underlying member method code with example usage: obj.init(w) and returns: momentum,  v1s,  v2s