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: Computes updated layer weights and optimizer slots for one training step.
  
  Args:
    step: Training step number.
    grads: Gradient values for this node (from back-propagation during a
        training step).
    weights: Current weight values for this node (i.e., layer weights).
    slots: Current slot values for this node.
    opt_params: Optimizer hyperparameters (e.g. learning rate, momentum),
        same across all nodes in the model.
  
  Returns:
    Tuple of (new_weights, new_slots), which the Trax runtime will use to
    update the model and optimizer within each training step.
docstr_tok: ["Computes","updated","layer","weights","and","optimizer","slots","for","one","training","step",".","Args",":","step",":","Training","step","number",".","grads",":","Gradient","values","for","this","node","(","from","back-propagation","during","a","training","step",")",".","weights",":","Current","weight","values","for","this","node","(","i.e.",",","layer","weights",")",".","slots",":","Current","slot","values","for","this","node",".","opt_params",":","Optimizer","hyperparameters","(","e.g",".","learning","rate",",","momentum",")",",","same","across","all","nodes","in","the","model",".","Returns",":","Tuple","of","(","new_weights",",","new_slots",")",",","which","the","Trax","runtime","will","use","to","update","the","model","and","optimizer","within","each","training","step","."]
code: def update(self, step, grads, weights, slots, opt_params):
      """Computes updated layer weights and optimizer slots for one training step.
  
      Args:
        step: Training step number.
        grads: Gradient values for this node (from back-propagation during a
            training step).
        weights: Current weight values for this node (i.e., layer weights).
        slots: Current slot values for this node.
        opt_params: Optimizer hyperparameters (e.g. learning rate, momentum),
            same across all nodes in the model.
  
      Returns:
        Tuple of (new_weights, new_slots), which the Trax runtime will use to
        update the model and optimizer within each training step.
      """
      raise NotImplementedError
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 3 for calling the Optimizer obj's underlying member method code with example usage: obj.update(step, grads, weights, slots, opt_params) without return types