code: def accuracy(params, batch):
      inputs, targets = batch
      target_class = jnp.argmax(targets, axis=1)
      predicted_class = jnp.argmax(predict(params, inputs), axis=1)
      return jnp.mean(predicted_class == target_class)
  
file_path: jax/examples/spmd_mnist_classifier_fromscratch.py
imports: from functools import partial
  import time
  import numpy
  import numpy.random
  import jax
  from jax import jit, grad, pmap
  from jax.scipy.special import logsumexp
  from jax.tree_util import tree_map
  from jax import lax
  import jax.numpy
  from examples import datasets
label_desc: Use image node_id 4 for calling a global function with example usage: accuracy(params, batch) and returns: jnp