class_docstr: Model provides an Estimator-like API similar to Keras.
code: def train_step(
      self: M,
      inputs: tp.Any,
      labels: tp.Mapping[str, tp.Any],
  ) -> TrainStepOutput[M]:
      model: M = self
      grads: M
      logs: types.Logs
      model: M = self
  
      if model.optimizer is None:
          raise types.MissingModule(
              "Trying to run default `train_step` on a Model with no `optimizer`, try overriding `train_step` or set `optimizer`"
          )
  
      params = model.parameters(self._is_trainable)
  
      grad_fn = jax.grad(self.loss_fn, has_aux=True)
      grads, (logs, model) = grad_fn(params, model, inputs, labels)
  
      model, grads = model.distributed_strategy.handle_model_and_grads(
          model, grads
      )
  
      assert model.optimizer is not None
  
      params = model.optimizer.update(grads, params)
      model = model.merge(params)
  
      return logs, model
  
imports: import typing
  import flax
  import jax
  import jax.numpy
  import numpy
  import treex
  from optax import GradientTransformation
  from elegy import types, utils
  from elegy.model.model_base import ModelBase
  from elegy.model.model_core import GradStepOutput, PredStepOutput, TestStepOutput, TrainStepOutput
label_desc: Use image node_id 10 for calling the Model obj's underlying member method code with example usage: obj.train_step(inputs, labels) and returns: logs,  model