class_docstr: Base class for optimizers that work hand in hand with Trax layers.
  
  To define an optimizer subclass, specify its behavior with respect to a
  single node in the network (e.g., a single dense layer):
  
    - `init`: how to create/initialize optimizer-internal parameters ("slots"),
        as a function of the node's weights.
    - `update`: how to use gradient information to update node weights and
        optimizer slots.
  
  The Trax runtime combines these node-local computations into layer weight
  updates and optimizer slot updates for the whole tree of layers in the model.
class_docstr_tok: ["Base","class","for","optimizers","that","work","hand","in","hand","with","Trax","layers",".","To","define","an","optimizer","subclass",",","specify","its","behavior","with","respect","to","a","single","node","in","the","network","(","e.g.",",","a","single","dense","layer",")",":","-","`","init","`",":","how","to","create\/initialize","optimizer-internal","parameters","(","``","slots","''",")",",","as","a","function","of","the","node","'s","weights",".","-","`","update","`",":","how","to","use","gradient","information","to","update","node","weights","and","optimizer","slots",".","The","Trax","runtime","combines","these","node-local","computations","into","layer","weight","updates","and","optimizer","slot","updates","for","the","whole","tree","of","layers","in","the","model","."]
docstr: Assembles node-local weight and slot updates for the full layer tree.
  
  Args:
    step: Current step number in the training process.
    grad_tree: Gradients for the entire model, in a tree that matches the
        model's layer structure.
    weight_tree: Current weights for the entire model, in a tree that matches
        the model's layer structure.
    slots: Optimizer slots.
    opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).
    store_slots: Boolean; if True, stores resulting slots in this object;
      when set to False, this becomes a pure function.
  
  Returns:
    Tuple `(weights, slots)`, where `weights` are the optimizer-updated
    weights for the whole model (in a tree matching the model's layer
    structure) and `slots` are the updated optimizer slot values.
docstr_tok: ["Assembles","node-local","weight","and","slot","updates","for","the","full","layer","tree",".","Args",":","step",":","Current","step","number","in","the","training","process",".","grad_tree",":","Gradients","for","the","entire","model",",","in","a","tree","that","matches","the","model","'s","layer","structure",".","weight_tree",":","Current","weights","for","the","entire","model",",","in","a","tree","that","matches","the","model","'s","layer","structure",".","slots",":","Optimizer","slots",".","opt_params",":","Optimizer","hyperparameters","(","e.g",".","learning","rate",",","momentum",")",".","store_slots",":","Boolean",";","if","True",",","stores","resulting","slots","in","this","object",";","when","set","to","False",",","this","becomes","a","pure","function",".","Returns",":","Tuple","`","(","weights",",","slots",")","`",",","where","`","weights","`","are","the","optimizer-updated","weights","for","the","whole","model","(","in","a","tree","matching","the","model","'s","layer","structure",")","and","`","slots","`","are","the","updated","optimizer","slot","values","."]
code: def tree_update(
      self,
      step,
      grad_tree,
      weight_tree,
      slots,
      opt_params,
      store_slots=True,
  ):
      """Assembles node-local weight and slot updates for the full layer tree.
  
      Args:
        step: Current step number in the training process.
        grad_tree: Gradients for the entire model, in a tree that matches the
            model's layer structure.
        weight_tree: Current weights for the entire model, in a tree that matches
            the model's layer structure.
        slots: Optimizer slots.
        opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).
        store_slots: Boolean; if True, stores resulting slots in this object;
          when set to False, this becomes a pure function.
  
      Returns:
        Tuple `(weights, slots)`, where `weights` are the optimizer-updated
        weights for the whole model (in a tree matching the model's layer
        structure) and `slots` are the updated optimizer slot values.
      """
      grads_flat = fastmath.tree_flatten(grad_tree)
      grads_norm = self._l2_norm(grads_flat)
      if self._clip_grad_norm is not None:
          max_norm = self._clip_grad_norm
          grads_flat = [
              jnp.where(
                  grads_norm
                  < max_norm,  # pylint: disable=g-complex-comprehension
                  g,
                  g * (max_norm / grads_norm),
              )
              for g in grads_flat
          ]
      weights_flat = fastmath.tree_flatten(weight_tree)
      weights_norm = self._l2_norm(weights_flat)
      updated_pairs = [
          self._update_and_check(step, grad, weight, slot, opt_params)
          for (grad, weight, slot) in zip(
              grads_flat, weights_flat, slots
          )
      ]
      new_weights_flat, slots = zip(*updated_pairs)
      new_weights, _ = fastmath.tree_unflatten(
          new_weights_flat, weight_tree
      )
      metrics = {"gradients_l2": grads_norm, "weights_l2": weights_norm}
      slots = tuple(slots)
      if store_slots:
          self.slots = slots
      return new_weights, slots, metrics
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 9 for calling the Optimizer obj's underlying member method code with example usage: obj.tree_update(step, grad_tree, weight_tree, slots, opt_params, store_slots) and returns: new_weights,  slots,  metrics