code: def test_step(
      self: M,
      inputs,
      labels,
  ) -> eg.TestStepOutput[M]:
      model: M = self
      # flatten + scale
      inputs = jnp.reshape(inputs, (inputs.shape[0], -1)) / 255
  
      # forward
      logits, model = model.pred_step(inputs)
  
      # crossentropy loss
      target = jax.nn.one_hot(labels["target"], self.features_out)
      loss = optax.softmax_cross_entropy(logits, target).mean()
  
      # metrics
      logs = dict(
          acc=jnp.mean(jnp.argmax(logits, axis=-1) == labels["target"]),
          loss=loss,
      )
  
      return loss, logs, model
  
file_path: elegy/examples/jax/linear_classifier_train_step.py
imports: import os
  import typing
  from datetime import datetime
  import jax
  import jax.numpy
  import numpy
  import optax
  import typer
  from datasets.load import load_dataset
  import elegy
label_desc: Use image node_id 4 for calling the Model obj's underlying member method code with example usage: obj.test_step(inputs, labels) and returns: loss,  logs,  model