import jax
import jax.numpy as jnp
from flax import linen as nn

Array = jax.Array


class Cifar10CNN(nn.Module):
    num_classes: int = 10
    channels: tuple[int, ...] = (64, 64, 128, 128, 256, 256)
    groups: int = 8

    @nn.compact
    def __call__(self, x: Array) -> Array:
        for i, c in enumerate(self.channels):
            x = nn.Conv(c, kernel_size=(3, 3), padding="SAME")(x)
            x = nn.GroupNorm(num_groups=min(self.groups, c))(x)
            x = nn.relu(x)
            if i % 2 == 1:
                x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = jnp.mean(x, axis=(1, 2))
        return nn.Dense(self.num_classes)(x)