code: def predict(model, images):
      input = images
      for w, b in model[:-1]:
          output = jnp.dot(input, w) + b
          input = jnp.tanh(output)
  
      last_w, last_b = model[-1]
      last_output = jnp.dot(input, last_w) + last_b
      return last_output - logsumexp(last_output, axis=1, keepdims=True)
  
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 2 for calling a global function with example usage: predict(model, images) and returns: unknown