code: def loss(params, batch):
      inputs, targets = batch
      preds = predict(params, inputs)
      return -jnp.mean(jnp.sum(preds * targets, axis=1))
  
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 3 for calling a global function with example usage: loss(params, batch) and returns: unknown