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 initializations into full-tree initialization.
  
  Args:
    weight_tree: Weights for an entire model, in a tree that matches the
        model's layer structure.
  
  Returns:
    Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer
    slot values and `opt_params` are optimizer hyperparameters (e.g.,
    learning rate, momentum).
docstr_tok: ["Assembles","node-local","initializations","into","full-tree","initialization",".","Args",":","weight_tree",":","Weights","for","an","entire","model",",","in","a","tree","that","matches","the","model","'s","layer","structure",".","Returns",":","Tuple","`","(","slots",",","opt_params",")","`",",","where","`","slots","`","are","the","initialized","optimizer","slot","values","and","`","opt_params","`","are","optimizer","hyperparameters","(","e.g.",",","learning","rate",",","momentum",")","."]
code: def tree_init(self, weight_tree):
      """Assembles node-local initializations into full-tree initialization.
  
      Args:
        weight_tree: Weights for an entire model, in a tree that matches the
            model's layer structure.
  
      Returns:
        Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer
        slot values and `opt_params` are optimizer hyperparameters (e.g.,
        learning rate, momentum).
      """
      self._slots = tuple(
          self.init(weight)
          for weight in fastmath.tree_flatten(weight_tree)
      )
      return (self._slots, self._init_opt_params)
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 8 for calling the Optimizer obj's underlying member method code with example usage: obj.tree_init(weight_tree) and returns: self,  self