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
Updates a single weight array and checks types.
trax/optimizers/base.py