code: def net_fn(images: jax.Array) -> jax.Array:
      """Standard LeNet-300-100 MLP network."""
      x = images.astype(jnp.float32) / 255.0
      mlp = hk.Sequential(
          [
              hk.Flatten(),
              hk.Linear(300),
              jax.nn.relu,
              hk.Linear(100),
              jax.nn.relu,
              hk.Linear(NUM_CLASSES),
          ]
      )
      return mlp(x)
  
file_path: dm-haiku/examples/mnist.py
imports: from collections.abc import Iterator
  from typing import NamedTuple
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: net_fn(images) and returns: mlp