class_docstr: Model provides an Estimator-like API similar to Keras.
code: def test_step(
      self: M,
      inputs: tp.Any,
      labels: tp.Mapping[str, tp.Any],
  ) -> TestStepOutput[M]:
      model: M = self
  
      if model.module is None:
          raise types.MissingModule(
              "Trying to run default `test_step` on a Model with no `loss_and_logs`, try overriding `test_step` or set `loss_and_logs`"
          )
  
      preds, model = model.pred_step(inputs)
      assert model.loss_and_logs is not None
  
      aux_losses = model.loss_logs()
      aux_metrics = model.metric_logs()
  
      extended_labels = {
          "inputs": inputs,
          "preds": preds,
          "model": model,
          "parameters": model.parameters(),
          "batch_stats": model.batch_stats(),
          "rngs": model.rngs(),
          "model_states": model.model_states(),
          "states": model.states(),
          "metric_logs": model.metric_logs(),
          "loss_logs": model.loss_logs(),
          "logs": model.logs(),
          **labels,
      }
  
      losses_kwargs = extended_labels
      metrics_kwargs = extended_labels
  
      # loss, *_ = model.loss_and_logs(
      #     **losses_kwargs,
      #     metrics_kwargs=metrics_kwargs,
      #     aux_losses=aux_losses,
      #     aux_metrics=aux_metrics,
      # )
  
      # TODO: could be refactored as reset_update in treex?
      batch_loss_and_logs = model.loss_and_logs.copy()
      batch_loss_and_logs.reset()
      batch_loss_and_logs.update(
          **losses_kwargs,
          metrics_kwargs=metrics_kwargs,
          aux_losses=aux_losses,
          aux_metrics=aux_metrics,
      )
      loss, *_ = batch_loss_and_logs.compute()
  
      batch_loss_and_logs = model.distributed_strategy.handle_metrics(
          batch_loss_and_logs
      )
  
      self.loss_and_logs = jax.tree_map(
          lambda a, b: a + b, self.loss_and_logs, batch_loss_and_logs
      )
  
      _, losses_logs, metrics_logs = model.loss_and_logs.compute()
  
      logs = {**losses_logs, **metrics_logs}
  
      return loss, 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 7 for calling the Model obj's underlying member method code with example usage: obj.test_step(inputs, labels) and returns: loss,  logs,  model