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: Sets initial hyperparameter values for this optimizer.
  
  Takes optimizer hyperparameters as keyword arguments. These values can
  change over time (training steps), e.g., for learning rate schedules.
  
  To expose subclass hyperparameters for gin configuration, override this
  constructor and use explicitly named keyword arguments. See
  `momentum.Momentum.__init__` for one such example.
  
  Args:
    learning_rate: Learning rate for the optimizer. This can change during
        training by means of a training rate schedule.
    clip_grad_norm: If specified, this scalar value is used to limit gradient
        size -- all gradient elements in a training step are treated as if
        they belonged to a single vector and then scaled back if needed so
        that such a vector's L2 norm does not exceed `clip_grad_norm`. If
        None, no clipping happens.
    **init_opt_params: Initial values of any additional optimizer parameters.
docstr_tok: ["Sets","initial","hyperparameter","values","for","this","optimizer",".","Takes","optimizer","hyperparameters","as","keyword","arguments",".","These","values","can","change","over","time","(","training","steps",")",",","e.g.",",","for","learning","rate","schedules",".","To","expose","subclass","hyperparameters","for","gin","configuration",",","override","this","constructor","and","use","explicitly","named","keyword","arguments",".","See","`","momentum.Momentum.__init__","`","for","one","such","example",".","Args",":","learning_rate",":","Learning","rate","for","the","optimizer",".","This","can","change","during","training","by","means","of","a","training","rate","schedule",".","clip_grad_norm",":","If","specified",",","this","scalar","value","is","used","to","limit","gradient","size","--","all","gradient","elements","in","a","training","step","are","treated","as","if","they","belonged","to","a","single","vector","and","then","scaled","back","if","needed","so","that","such","a","vector","'s","L2","norm","does","not","exceed","`","clip_grad_norm","`",".","If","None",",","no","clipping","happens",".","*","*","init_opt_params",":","Initial","values","of","any","additional","optimizer","parameters","."]
code: def __init__(
      self, learning_rate=0.01, clip_grad_norm=None, **init_opt_params
  ):
      """Sets initial hyperparameter values for this optimizer.
  
      Takes optimizer hyperparameters as keyword arguments. These values can
      change over time (training steps), e.g., for learning rate schedules.
  
      To expose subclass hyperparameters for gin configuration, override this
      constructor and use explicitly named keyword arguments. See
      `momentum.Momentum.__init__` for one such example.
  
      Args:
        learning_rate: Learning rate for the optimizer. This can change during
            training by means of a training rate schedule.
        clip_grad_norm: If specified, this scalar value is used to limit gradient
            size -- all gradient elements in a training step are treated as if
            they belonged to a single vector and then scaled back if needed so
            that such a vector's L2 norm does not exceed `clip_grad_norm`. If
            None, no clipping happens.
        **init_opt_params: Initial values of any additional optimizer parameters.
      """
      init_opt_params["learning_rate"] = learning_rate
      self._init_opt_params = {
          name: jnp.array(value)
          for (name, value) in init_opt_params.items()
      }
      self._slots = None
      # Gradient clipping happens with respect to the norm of the whole gradient
      # tree, so it is not passed to single-slot updates, but done in this class
      # for the whole gradient tree.
      self._clip_grad_norm = clip_grad_norm
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 1 to create a new Optimizer object  with example: obj = Optimizer(learning_rate, clip_grad_norm)