Add a second variant of MNIST CNN

tests/jax/models/mnist/cnn/dropout/model_implementation.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from flax import linen as nn
+
+
+class MNISTCNNDropoutModel(nn.Module):
+    @nn.compact
+    def __call__(self, x, *, train: bool):
+        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
+        x = nn.relu(x)
+
+        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
+        x = nn.relu(x)
+        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
+
+        x = nn.Dropout(rate=0.25)(x, deterministic=not train)
+
+        x = x.reshape((x.shape[0], -1))
+
+        x = nn.Dense(features=128)(x)
+        x = nn.relu(x)
+        x = nn.Dropout(rate=0.5)(x, deterministic=not train)
+
+        x = nn.Dense(features=10)(x)
+        x = nn.softmax(x)
+
+        return x

tests/jax/models/mnist/cnn/dropout/test_mnist_cnn_dropout.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import pytest
+from infra import RunMode
+
+from ..tester import MNISTCNNTester
+from .model_implementation import MNISTCNNDropoutModel
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> MNISTCNNTester:
+    return MNISTCNNTester(MNISTCNNDropoutModel)
+
+
+@pytest.fixture
+def training_tester() -> MNISTCNNTester:
+    return MNISTCNNTester(MNISTCNNDropoutModel, RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.xfail(
+    reason="error: failed to legalize unresolved materialization from () to 'tensor<1xf32>' that remained live after conversion"
+)
+def test_mnist_cnn_dropout_inference(
+    inference_tester: MNISTCNNTester,
+):
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_mnist_cnn_nodropout_training(
+    training_tester: MNISTCNNTester,
+):
+    training_tester.test()

tests/jax/models/mnist/cnn/model_implementation.py → tests/jax/models/mnist/cnn/nodropout/model_implementation.py

@@ -2,11 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import jax
 from flax import linen as nn
 
 
-class MNISTCNNModel(nn.Module):
+class MNISTCNNNoDropoutModel(nn.Module):
     @nn.compact
     def __call__(self, x, *, train: bool):
         x = nn.Conv(features=32, kernel_size=(3, 3), padding="SAME")(x)

tests/jax/models/mnist/cnn/nodropout/test_mnist_cnn_nodropout.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import pytest
+from infra import RunMode
+
+from ..tester import MNISTCNNTester
+from .model_implementation import MNISTCNNNoDropoutModel
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> MNISTCNNTester:
+    return MNISTCNNTester(MNISTCNNNoDropoutModel)
+
+
+@pytest.fixture
+def training_tester() -> MNISTCNNTester:
+    return MNISTCNNTester(MNISTCNNNoDropoutModel, RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.skip(
+    reason='void mlir::OperationConverter::finalize(mlir::ConversionPatternRewriter &): Assertion `newValue && "replacement value not found"\' failed.'
+)  # This is a segfault, marking it as xfail would bring down the whole test suite
+def test_mnist_cnn_nodropout_inference(
+    inference_tester: MNISTCNNTester,
+):
+    inference_tester.test()
+
+
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_mnist_cnn_nodropout_training(
+    training_tester: MNISTCNNTester,
+):
+    training_tester.test()

tests/jax/models/mnist/cnn/test_mnist_cnn.py → tests/jax/models/mnist/cnn/tester.py

@@ -6,19 +6,20 @@
 
 import jax
 import jax.numpy as jnp
-import pytest
 from flax import linen as nn
-from infra import ModelTester, RunMode
-
-from tests.jax.models.mnist.cnn.model_implementation import MNISTCNNModel
+from infra import ModelTester
 
 
 class MNISTCNNTester(ModelTester):
     """Tester for MNIST CNN model."""
 
+    def __init__(self, cls):
+        self._model_class = cls
+        super().__init__()
+
     # @override
     def _get_model(self) -> nn.Module:
-        return MNISTCNNModel()
+        return self._model_class()
 
     # @override
     def _get_forward_method_name(self) -> str:
@@ -47,39 +48,3 @@ def _get_forward_method_kwargs(self) -> Dict[str, jax.Array]:
     # @override
     def _get_static_argnames(self):
         return ["train"]
-
-
-# ----- Fixtures -----
-
-
-@pytest.fixture
-def inference_tester() -> MNISTCNNTester:
-    return MNISTCNNTester()
-
-
-@pytest.fixture
-def training_tester() -> MNISTCNNTester:
-    return MNISTCNNTester(RunMode.TRAINING)
-
-
-# ----- Tests -----
-
-
-@pytest.mark.skip(
-    reason='void mlir::OperationConverter::finalize(mlir::ConversionPatternRewriter &): Assertion `newValue && "replacement value not found"\' failed.'
-)  # This is a segfault, marking it as xfail would bring down the whole test suite
-def test_mnist_cnn_inference(
-    inference_tester: MNISTCNNTester,
-):
-    inference_tester.test()
-
-
-@pytest.mark.skip(reason="Support for training not implemented")
-def test_mnist_cnn_training(
-    training_tester: MNISTCNNTester,
-):
-    training_tester.test()
-
-
-if __name__ == "__main__":
-    MNISTCNNTester().test()