code: def eval_step(
      state: nnx.TrainState[MLP], inputs: jax.Array, labels: jax.Array
  ):
      logits, _ = state.apply("params")(inputs)
      loss = jnp.mean(
          optax.softmax_cross_entropy_with_integer_labels(
              logits, labels
          )
      )
      acc = jnp.mean(jnp.argmax(logits, axis=-1) == labels)
      return {"loss": loss, "accuracy": acc}
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  from datasets import load_dataset
  from flax.experimental import nnx
label_desc: Use image node_id 9 for calling a global function with example usage: eval_step(state, inputs, labels) and returns: dict