code: def forward(x: jnp.ndarray, training: bool):
      # Normalize input
      x = x.astype(jnp.float32) / 255.0
  
      # Block 1
      x = hk.Conv2D(32, [3, 3], stride=[2, 2])(x)
      x = hk.BatchNorm(
          create_scale=True, create_offset=True, decay_rate=0.99
      )(x, is_training=training)
      x = hk.dropout(hk.next_rng_key(), 0.05, x)
      x = jax.nn.relu(x)
  
      # Block 2
      x = hk.Conv2D(64, [3, 3], stride=[2, 2])(x)
      x = hk.BatchNorm(
          create_scale=True, create_offset=True, decay_rate=0.99
      )(x, is_training=training)
      x = hk.dropout(hk.next_rng_key(), 0.1, x)
      x = jax.nn.relu(x)
  
      # Block 3
      x = hk.Conv2D(128, [3, 3], stride=[2, 2])(x)
  
      # GlobalAveragePooling2D
      x = x.mean(axis=(1, 2))
  
      # Classification layer
      x = hk.Linear(10)(x)
  
      return x
  
file_path: elegy/examples/haiku/mnist_conv.py
imports: import os
  import typing
  from datetime import datetime
  import datasets
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  import typer
  from datasets import load_dataset
  import elegy
label_desc: Use image node_id 1 for calling a global function with example usage: forward(x, training) and returns: x