from functools import partial

class CNN(nnx.Module):
    """Convolutional Neural Network for MNIST."""
    
    def __init__(self, *, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
        self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
        self.linear2 = nnx.Linear(256, 10, rngs=rngs)
        self.pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
    
    def __call__(self, x):
        x = nnx.relu(self.conv1(x))
        x = self.pool(x)
        x = nnx.relu(self.conv2(x))
        x = self.pool(x)
        x = x.reshape(x.shape[0], -1)  # Flatten
        x = nnx.relu(self.linear1(x))
        x = self.linear2(x)
        return x

        __
        __