code: def sm3(step_size, momentum=0.9):
      """Construct optimizer triple for SM3.
  
      Memory-Efficient Adaptive Optimization for Large-Scale Learning.
      https://arxiv.org/abs/1901.11150
  
      Args:
        step_size: positive scalar, or a callable representing a step size schedule
          that maps the iteration index to a positive scalar.
        momentum: optional, a positive scalar value for momentum
  
      Returns:
        An (init_fun, update_fun, get_params) triple.
      """
      step_size = make_schedule(step_size)
  
      def splice(seq, i, x):
          lst = list(seq)
          lst[i : i + 1] = x
          return lst
  
      def broadcast_into(ndim, x, axis):
          idx = splice([None] * ndim, axis, [slice(None)])
          return x[tuple(idx)]
  
      def init(x0):
          x_shape = x0.shape
          x0 = jnp.atleast_1d(x0)
          vs = [jnp.zeros(sz, dtype=x0.dtype) for sz in x0.shape]
          return x0, jnp.zeros_like(x0), vs, x_shape
  
      def update(i, g, state):
          x, m, vs, x_shape = state
          vs = [broadcast_into(g.ndim, v, i) for i, v in enumerate(vs)]
          accum = functools.reduce(jnp.minimum, vs) + jnp.square(g)
          accum_inv_sqrt = jnp.where(
              accum > 0, 1.0 / jnp.sqrt(accum), 0
          )
          m = (1.0 - momentum) * (g * accum_inv_sqrt) + momentum * m
          x = x - step_size(i) * m
          vs = [
              accum.max(splice(range(x.ndim), j, []))
              for j in range(x.ndim)
          ]
          return x, m, vs, x_shape
  
      def get_params(state):
          x, _, _, x_shape = state
          return x.reshape(x_shape)
  
      return init, update, get_params
  
imports: from __future__ import annotations
  from typing import Any, Callable, NamedTuple
  from collections import namedtuple
  import functools
  from functools import partial
  import jax.numpy
  from jax._src.util import safe_zip, safe_map, unzip2
  from jax import tree_util
  from jax.tree_util import tree_map, tree_flatten, tree_unflatten, register_pytree_node
label_desc: Use image node_id 10 for calling a global function with example usage: sm3(step_size, momentum) and returns: init,  update,  get_params, lst, x, x0,  jnp,  vs,  x_shape, x,  m,  vs,  x_shape, x