code: def predict(params, inputs):
      activations = inputs
      for w, b in params[:-1]:
          outputs = jnp.dot(activations, w) + b
          activations = jnp.tanh(outputs)
  
      final_w, final_b = params[-1]
      logits = jnp.dot(activations, final_w) + final_b
      return logits - logsumexp(logits, axis=1, keepdims=True)
  
file_path: jax/examples/mnist_classifier_fromscratch.py
imports: import time
  import numpy.random
  from jax import jit, grad
  from jax.scipy.special import logsumexp
  import jax.numpy
  from examples import datasets
label_desc: Use image node_id 2 for calling a global function with example usage: predict(params, inputs) and returns: unknown