code: def net_fn(batch: Batch) -> jax.Array:
      """Standard LeNet-300-100 MLP network."""
      x = batch["image"].astype(jnp.float32) / 255.0
      mlp = hk.Sequential(
          [
              hk.Flatten(),
              hk.Linear(300),
              jax.nn.relu,
              hk.Linear(100),
              jax.nn.relu,
              hk.Linear(10),
          ]
      )
      return mlp(x)
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 8 for calling a global function with example usage: net_fn(batch) and returns: mlp