human
Give me a the Python program you requested given the following:
2. Define network
Create a convolutional neural network with the Linen API by subclassing
Module.
Because the architecture in this example is relatively simple—you're just
stacking layers—you can define the inlined submodules directly within the
__call__ method and wrap it with the
@compact
decorator.
assistant
#Here is the Python program you requested.
class CNN(nn.Module):
  """A simple CNN model."""

  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x