code: def loss(model, batch):
      predicted_labels = predict(model, batch["images"])
      return -jnp.mean(
          jnp.sum(predicted_labels * batch["labels"], axis=1)
      )
  
file_path: DALI/docs/examples/frameworks/jax/model.py
imports: from jax import jit, grad, lax, pmap
  from functools import partial
  import jax.numpy
  from jax.scipy.special import logsumexp
  import numpy.random
label_desc: Use image node_id 3 for calling a global function with example usage: loss(model, batch) and returns: unknown