code: def accuracy(model, iterator):
      correct_predictions_num = 0
      for batch in iterator:
          images = batch["images"]
          labels = batch["labels"]
  
          predicted_class = jnp.argmax(predict(model, images), axis=1)
          correct_predictions_num = correct_predictions_num + jnp.sum(
              predicted_class == labels.ravel()
          )
  
      return correct_predictions_num / iterator.size
  
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 4 for calling a global function with example usage: accuracy(model, iterator) and returns: unknown