Method __call__

tests/test_flax.py:23–35

(self, x)

    # cuDNN 9.19 (pulled in by torch 2.11) dropped the Pascal (sm_60) kernels from
    # libcudnn_ops/cnn/adv and ships PTX for sm_121 only, so nothing can be JIT'd
    # down to sm_60 either. Every cuDNN convolution engine fails on P100 with
    # CUDNN_STATUS_EXECUTION_FAILED. Not fixable here: cuDNN 9.19 is a hard
    # requirement of torch 2.11, which comes from the Colab base image.
    @p100_exempt
    def test_cnn(self):
        class CNN(nn.Module):
            @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))