class_docstr: Simple train state for the common case with a single Optax optimizer.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> from flax.training.train_state import TrainState
    >>> import jax, jax.numpy as jnp
    >>> import optax
  
    >>> x = jnp.ones((1, 2))
    >>> y = jnp.ones((1, 2))
    >>> model = nn.Dense(2)
    >>> variables = model.init(jax.random.key(0), x)
    >>> tx = optax.adam(1e-3)
  
    >>> state = TrainState.create(
    ...     apply_fn=model.apply,
    ...     params=variables['params'],
    ...     tx=tx)
  
    >>> def loss_fn(params, x, y):
    ...   predictions = state.apply_fn({'params': params}, x)
    ...   loss = optax.l2_loss(predictions=predictions, targets=y).mean()
    ...   return loss
    >>> loss_fn(state.params, x, y)
    Array(3.3514676, dtype=float32)
  
    >>> grads = jax.grad(loss_fn)(state.params, x, y)
    >>> state = state.apply_gradients(grads=grads)
    >>> loss_fn(state.params, x, y)
    Array(3.343844, dtype=float32)
  
  Note that you can easily extend this dataclass by subclassing it for storing
  additional data (e.g. additional variable collections).
  
  For more exotic usecases (e.g. multiple optimizers) it's probably best to
  fork the class and modify it.
  
  Args:
    step: Counter starts at 0 and is incremented by every call to
      ``.apply_gradients()``.
    apply_fn: Usually set to ``model.apply()``. Kept in this dataclass for
      convenience to have a shorter params list for the ``train_step()`` function
      in your training loop.
    params: The parameters to be updated by ``tx`` and used by ``apply_fn``.
    tx: An Optax gradient transformation.
    opt_state: The state for ``tx``.
docstr: Updates ``step``, ``params``, ``opt_state`` and ``**kwargs`` in return value.
  
  Note that internally this function calls ``.tx.update()`` followed by a call
  to ``optax.apply_updates()`` to update ``params`` and ``opt_state``.
  
  Args:
    grads: Gradients that have the same pytree structure as ``.params``.
    **kwargs: Additional dataclass attributes that should be ``.replace()``-ed.
  
  Returns:
    An updated instance of ``self`` with ``step`` incremented by one, ``params``
    and ``opt_state`` updated by applying ``grads``, and additional attributes
    replaced as specified by ``kwargs``.
docstr_tok: ["Updates","``","step","``",",","``","params","``",",","``","opt_state","``","and","``","*","*","kwargs","``","in","return","value",".","Note","that","internally","this","function","calls","``",".tx.update","(",")","``","followed","by","a","call","to","``","optax.apply_updates","(",")","``","to","update","``","params","``","and","``","opt_state","``",".","Args",":","grads",":","Gradients","that","have","the","same","pytree","structure","as","``",".params","``",".","*","*","kwargs",":","Additional","dataclass","attributes","that","should","be","``",".replace","(",")","``","-ed",".","Returns",":","An","updated","instance","of","``","self","``","with","``","step","``","incremented","by","one",",","``","params","``","and","``","opt_state","``","updated","by","applying","``","grads","``",",","and","additional","attributes","replaced","as","specified","by","``","kwargs","``","."]
code: def apply_gradients(self, *, grads, **kwargs):
      """Updates ``step``, ``params``, ``opt_state`` and ``**kwargs`` in return value.
  
      Note that internally this function calls ``.tx.update()`` followed by a call
      to ``optax.apply_updates()`` to update ``params`` and ``opt_state``.
  
      Args:
        grads: Gradients that have the same pytree structure as ``.params``.
        **kwargs: Additional dataclass attributes that should be ``.replace()``-ed.
  
      Returns:
        An updated instance of ``self`` with ``step`` incremented by one, ``params``
        and ``opt_state`` updated by applying ``grads``, and additional attributes
        replaced as specified by ``kwargs``.
      """
      if OVERWRITE_WITH_GRADIENT in grads:
          grads_with_opt = grads["params"]
          params_with_opt = self.params["params"]
      else:
          grads_with_opt = grads
          params_with_opt = self.params
  
      updates, new_opt_state = self.tx.update(
          grads_with_opt, self.opt_state, params_with_opt
      )
      new_params_with_opt = optax.apply_updates(
          params_with_opt, updates
      )
  
      # As implied by the OWG name, the gradients are used directly to update the
      # parameters.
      if OVERWRITE_WITH_GRADIENT in grads:
          new_params = {
              "params": new_params_with_opt,
              OVERWRITE_WITH_GRADIENT: grads[OVERWRITE_WITH_GRADIENT],
          }
      else:
          new_params = new_params_with_opt
      return self.replace(
          step=self.step + 1,
          params=new_params,
          opt_state=new_opt_state,
          **kwargs,
      )
  
imports: from typing import Any, Callable
  import optax
  from flax import core, struct
  from flax.linen.fp8_ops import OVERWRITE_WITH_GRADIENT
label_desc: Use image node_id 1 for calling the TrainState obj's underlying member method code with example usage: obj.apply_gradients() and returns: self