class_docstr: Model provides an Estimator-like API similar to Keras.
code: def loss_fn(
      params: M,
      model: M,
      inputs: tp.Any,
      labels: tp.Mapping[str, tp.Any],
  ) -> tp.Tuple[jnp.ndarray, tp.Tuple[types.Logs, M]]:
      model = model.merge(params)
      loss, logs, model = model.test_step(inputs, labels)
      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 8 for calling the Model obj's underlying member method code with example usage: obj.loss_fn(params, model, inputs, labels) and returns: loss,  logs,  model