class SafeCNN(nnx.Module):
    def __init__(self, *, rngs: nnx.Rngs):
        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
        self.linear = nnx.Linear(21632, 10, rngs=rngs)
    
    def __call__(self, x):
        x = self.conv1(x)
        x = nnx.relu(x)
        
        # Catch problems right where they happen
        chex.assert_tree_all_finite(x)
        
        x = x.reshape(x.shape[0], -1)
        return self.linear(x)

        __
        __