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: Updates a single weight array and checks types.
docstr_tok: ["Updates","a","single","weight","array","and","checks","types","."]
code: def _update_and_check(self, step, grads, weights, slots, opt_params):
      """Updates a single weight array and checks types."""
      new_weights, new_slots = self.update(
          step, grads, weights, slots, opt_params
      )
      if isinstance(weights, jnp.ndarray):
          if not isinstance(new_weights, jnp.ndarray):
              raise ValueError(
                  f"New weight values should be of type jnp.ndarray or a subclass; "
                  f"instead got {type(new_weights)}."
              )
          if new_weights.dtype != weights.dtype:
              raise ValueError(
                  f"New weight values dtype ({new_weights.dtype}) does not match "
                  f"the old one ({weights.dtype})."
              )
      return new_weights, new_slots
  
file_path: trax/trax/optimizers/base.py
imports: from trax import fastmath
  from trax.fastmath import numpy
label_desc: Use image node_id 11 for calling the Optimizer obj's underlying member method code with example usage: obj._update_and_check(step, grads, weights, slots, opt_params) and returns: new_weights,  new_slots