Repository: tenstorrent/tt-xla
Pull request #361: Add alexnet model implementation and tests
Author: login: mrakitaTT
Created: 2025-03-16T23:40:22Z | Merged: 2025-03-17T18:10:54Z

### Ticket
/

### Problem description
AlexNet model wasn't brought up.

### What's changed
Added alexnet model implementation and tests. Fixed few small things in mnist tests.

### Checklist
- [x] New/Existing tests provide coverage for changes



Diff:
diff --git a/tests/jax/models/alexnet/__init__.py b/tests/jax/models/alexnet/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/jax/models/alexnet/model_implementation.py b/tests/jax/models/alexnet/model_implementation.py
new file mode 100644
index 0000000000..8c958a77b9
--- /dev/null
+++ b/tests/jax/models/alexnet/model_implementation.py
@@ -0,0 +1,130 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Optional, Union
+
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+
+
+class AlexNetModel(nn.Module):
+    """
+    JAX implementation of the AlexNet model originally introduced by Alex Krizhevsky
+    in papers:
+      - "ImageNet Classification with Deep Convolutional Neural Networks"
+      - "One weird trick for parallelizing convolutional neural networks"
+    """
+
+    param_dtype: Optional[
+        Union[str, type[Any], jnp.dtype, jax._src.typing.SupportsDType, Any]
+    ] = jnp.bfloat16
+
+    @nn.compact
+    def __call__(self, x: jax.Array, *, train: bool) -> jax.Array:
+        # First feature extraction layer
+        x = nn.Conv(
+            features=64,
+            kernel_size=(11, 11),
+            strides=(4, 4),
+            padding=0,
+            param_dtype=self.param_dtype,
+        )(x)
+        x = nn.relu(x)
+        x = LocalResponseNormalization()(x)
+        x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2))
+
+        # Second feature extraction layer
+        x = nn.Conv(
+            features=192,
+            kernel_size=(5, 5),
+            strides=(1, 1),
+            padding=2,
+            param_dtype=self.param_dtype,
+        )(x)
+        x = nn.relu(x)
+        x = LocalResponseNormalization()(x)
+        x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2))
+
+        # Third feature extraction layer
+        x = nn.Conv(
+            features=384,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            padding=1,
+            param_dtype=self.param_dtype,
+        )(x)
+        x = nn.relu(x)
+
+        # Fourth feature extraction layer
+        x = nn.Conv(
+            features=256,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            padding=1,
+            param_dtype=self.param_dtype,
+        )(x)
+        x = nn.relu(x)
+
+        # Fifth feature extraction layer
+        x = nn.Conv(
+            features=256,
+            kernel_size=(3, 3),
+            strides=(1, 1),
+            padding=1,
+            param_dtype=self.param_dtype,
+        )(x)
+        x = nn.relu(x)
+        x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2))
+
+        # First classifier layer
+        x = nn.Dense(features=4096, param_dtype=self.param_dtype)(x)
+        x = nn.relu(x)
+        x = nn.Dropout(rate=0.5)(x, deterministic=not train)
+
+        # Second classifier layer
+        x = nn.Dense(features=4096, param_dtype=self.param_dtype)(x)
+        x = nn.relu(x)
+        x = nn.Dropout(rate=0.5)(x, deterministic=not train)
+
+        # Third classifier layer
+        x = nn.Dense(features=1000, param_dtype=self.param_dtype)(x)
+        x = nn.softmax(x)
+
+        return x
+
+
+class LocalResponseNormalization(nn.Module):
+    """Local response normalization layer per original paper implementation."""
+
+    k: int = 2
+    n: int = 5
+    alpha: float = 1e-4
+    beta: float = 0.75
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        # Input is expected in (B, H, W, C) format
+        num_channels = x.shape[-1]
+        padded_x = jnp.pad(
+            x, pad_width=[(0, 0), (0, 0), (0, 0), (self.n // 2, self.n // 2)]
+        )
+
+        def _apply_per_channel(c):
+            window_sq = (
+                jax.lax.dynamic_slice_in_dim(
+                    operand=padded_x,
+                    start_index=c,
+                    slice_size=self.n,
+                    axis=3,
+                )
+            ) ** 2
+            window_sq_sum = (
+                jnp.einsum("bhwc->bhw", window_sq) * self.alpha + self.k
+            ) ** self.beta
+            return x[:, :, :, num_channels] / window_sq_sum
+
+        return jax.vmap(_apply_per_channel, in_axes=0, out_axes=3)(
+            jnp.arange(num_channels)
+        )
diff --git a/tests/jax/models/alexnet/test_alexnet.py b/tests/jax/models/alexnet/test_alexnet.py
new file mode 100644
index 0000000000..c74c4e171e
--- /dev/null
+++ b/tests/jax/models/alexnet/test_alexnet.py
@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Dict, Sequence
+
+import jax
+import jax.numpy as jnp
+import pytest
+from flax import linen as nn
+from infra import ComparisonConfig, Framework, ModelTester, RunMode
+
+from tests.utils import (
+    BringupStatus,
+    Category,
+    ModelGroup,
+    ModelSource,
+    ModelTask,
+    build_model_name,
+    failed_ttmlir_compilation,
+)
+
+from .model_implementation import AlexNetModel
+
+MODEL_NAME = build_model_name(
+    Framework.JAX,
+    "alexnet",
+    None,
+    ModelTask.CV_IMAGE_CLS,
+    ModelSource.CUSTOM,
+)
+
+
+# ----- Tester -----
+
+
+class AlexNetTester(ModelTester):
+    """Tester for AlexNet CNN model."""
+
+    def __init__(
+        self,
+        model_class,
+        comparison_config: ComparisonConfig = ComparisonConfig(),
+        run_mode: RunMode = RunMode.INFERENCE,
+    ) -> None:
+        self._model_class = model_class
+        super().__init__(comparison_config, run_mode)
+
+    # @override
+    def _get_model(self) -> nn.Module:
+        return self._model_class()
+
+    # @override
+    def _get_forward_method_name(self) -> str:
+        return "apply"
+
+    # @override
+    def _get_input_activations(self) -> Sequence[jax.Array]:
+        prng_key = jax.random.PRNGKey(23)
+        img = jax.random.randint(
+            key=prng_key,
+            # B, H, W, C
+            shape=(4, 224, 224, 3),
+            # In the original paper inputs are normalized with individual channel
+            # values learned from training set.
+            minval=-128,
+            maxval=128,
+        )
+        return img
+
+    # @override
+    def _get_forward_method_args(self):
+        inp = self._get_input_activations()
+
+        # Example of flax.linen convention of first instatiating a model object
+        # and then later calling init to generate a set of initial tensors (parameters
+        # and maybe some extra state). Parameters are not stored with the models
+        # themselves, they are provided together with inputs to the forward method.
+        parameters = self._model.init(jax.random.PRNGKey(42), inp, train=False)
+
+        return [parameters, inp]
+
+    # @override
+    def _get_forward_method_kwargs(self) -> Dict[str, jax.Array]:
+        return {"train": False}
+
+    # @override
+    def _get_static_argnames(self):
+        return ["train"]
+
+
+# ----- Fixtures -----
+
+
+@pytest.fixture
+def inference_tester() -> AlexNetTester:
+    return AlexNetTester(AlexNetModel)
+
+
+@pytest.fixture
+def training_tester() -> AlexNetTester:
+    return AlexNetTester(AlexNetModel, run_mode=RunMode.TRAINING)
+
+
+# ----- Tests -----
+
+
+@pytest.mark.model_test
+@pytest.mark.record_test_properties(
+    category=Category.MODEL_TEST,
+    model_name=MODEL_NAME,
+    model_group=ModelGroup.GENERALITY,
+    run_mode=RunMode.INFERENCE,
+    bringup_status=BringupStatus.FAILED_TTMLIR_COMPILATION,
+)
+@pytest.mark.xfail(
+    reason=failed_ttmlir_compilation(
+        "error: failed to legalize operation 'stablehlo.pad'"
+        "(https://github.com/tenstorrent/tt-xla/issues/360)"
+    )
+)
+def test_alexnet_inference(inference_tester: AlexNetTester):
+    inference_tester.test()
+
+
+@pytest.mark.nightly
+@pytest.mark.record_test_properties(
+    category=Category.MODEL_TEST,
+    model_name=MODEL_NAME,
+    model_group=ModelGroup.GENERALITY,
+    run_mode=RunMode.TRAINING,
+)
+@pytest.mark.skip(reason="Support for training not implemented")
+def test_alexnet_training(training_tester: AlexNetTester):
+    training_tester.test()
diff --git a/tests/jax/models/mnist/cnn/dropout/test_mnist_cnn_dropout.py b/tests/jax/models/mnist/cnn/dropout/test_mnist_cnn_dropout.py
index 8c3b9a68f7..0f77a299b6 100644
--- a/tests/jax/models/mnist/cnn/dropout/test_mnist_cnn_dropout.py
+++ b/tests/jax/models/mnist/cnn/dropout/test_mnist_cnn_dropout.py
@@ -18,8 +18,6 @@
 from ..tester import MNISTCNNTester
 from .model_implementation import MNISTCNNDropoutModel
 
-# ----- Fixtures -----
-
 MODEL_NAME = build_model_name(
     Framework.JAX,
     "mnist",
@@ -29,6 +27,9 @@
 )
 
 
+# ----- Fixtures -----
+
+
 @pytest.fixture
 def inference_tester() -> MNISTCNNTester:
     return MNISTCNNTester(MNISTCNNDropoutModel)
@@ -36,7 +37,7 @@ def inference_tester() -> MNISTCNNTester:
 
 @pytest.fixture
 def training_tester() -> MNISTCNNTester:
-    return MNISTCNNTester(MNISTCNNDropoutModel, RunMode.TRAINING)
+    return MNISTCNNTester(MNISTCNNDropoutModel, run_mode=RunMode.TRAINING)
 
 
 # ----- Tests -----
diff --git a/tests/jax/models/mnist/cnn/nodropout/test_mnist_cnn_nodropout.py b/tests/jax/models/mnist/cnn/nodropout/test_mnist_cnn_nodropout.py
index 57b9c720c3..6299b84830 100644
--- a/tests/jax/models/mnist/cnn/nodropout/test_mnist_cnn_nodropout.py
+++ b/tests/jax/models/mnist/cnn/nodropout/test_mnist_cnn_nodropout.py
@@ -18,8 +18,6 @@
 from ..tester import MNISTCNNTester
 from .model_implementation import MNISTCNNNoDropoutModel
 
-# ----- Fixtures -----
-
 MODEL_NAME = build_model_name(
     Framework.JAX,
     "mnist",
@@ -29,6 +27,9 @@
 )
 
 
+# ----- Fixtures -----
+
+
 @pytest.fixture
 def inference_tester() -> MNISTCNNTester:
     return MNISTCNNTester(MNISTCNNNoDropoutModel)
@@ -36,7 +37,7 @@ def inference_tester() -> MNISTCNNTester:
 
 @pytest.fixture
 def training_tester() -> MNISTCNNTester:
-    return MNISTCNNTester(MNISTCNNNoDropoutModel, RunMode.TRAINING)
+    return MNISTCNNTester(MNISTCNNNoDropoutModel, run_mode=RunMode.TRAINING)
 
 
 # ----- Tests -----
diff --git a/tests/jax/models/mnist/cnn/tester.py b/tests/jax/models/mnist/cnn/tester.py
index 8bf2f8be23..f008e8547d 100644
--- a/tests/jax/models/mnist/cnn/tester.py
+++ b/tests/jax/models/mnist/cnn/tester.py
@@ -7,15 +7,20 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from infra import ModelTester
+from infra import ComparisonConfig, ModelTester, RunMode
 
 
 class MNISTCNNTester(ModelTester):
     """Tester for MNIST CNN model."""
 
-    def __init__(self, cls):
-        self._model_class = cls
-        super().__init__()
+    def __init__(
+        self,
+        model_class,
+        comparison_config: ComparisonConfig = ComparisonConfig(),
+        run_mode: RunMode = RunMode.INFERENCE,
+    ) -> None:
+        self._model_class = model_class
+        super().__init__(comparison_config, run_mode)
 
     # @override
     def _get_model(self) -> nn.Module:
@@ -36,7 +41,9 @@ def _get_forward_method_args(self):
         inp = self._get_input_activations()
 
         # Example of flax.linen convention of first instatiating a model object
-        # and then later calling init to generate a set of initial tensors(parameters and maybe some extra state)
+        # and then later calling init to generate a set of initial tensors (parameters
+        # and maybe some extra state). Parameters are not stored with the models
+        # themselves, they are provided together with inputs to the forward method.
         parameters = self._model.init(jax.random.PRNGKey(42), inp, train=False)
 
         return [parameters, inp]
diff --git a/tests/jax/models/mnist/mlp/test_mnist_mlp.py b/tests/jax/models/mnist/mlp/test_mnist_mlp.py
index 7b2050c9f6..24ad7c8800 100644
--- a/tests/jax/models/mnist/mlp/test_mnist_mlp.py
+++ b/tests/jax/models/mnist/mlp/test_mnist_mlp.py
@@ -20,6 +20,17 @@
 
 from .model_implementation import MNISTMLPModel
 
+MODEL_NAME = build_model_name(
+    Framework.JAX,
+    "mnist",
+    "mlp",
+    ModelTask.CV_IMAGE_CLS,
+    ModelSource.CUSTOM,
+)
+
+
+# ----- Tester -----
+
 
 class MNISTMLPTester(ModelTester):
     """Tester for MNIST MLP model."""
@@ -60,15 +71,6 @@ def _get_forward_method_args(self):
 # ----- Fixtures -----
 
 
-MODEL_NAME = build_model_name(
-    Framework.JAX,
-    "mnist",
-    "mlp",
-    ModelTask.CV_IMAGE_CLS,
-    ModelSource.CUSTOM,
-)
-
-
 @pytest.fixture
 def inference_tester(request) -> MNISTMLPTester:
     return MNISTMLPTester(request.param)


Reviews:
user: login: kmitrovicTT
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T10:38:05Z
user: login: mrakitaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T10:42:05Z
user: login: sgligorijevicTT
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T11:20:33Z

Review comments:
diff_hunk: @@ -7,15 +7,20 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from infra import ModelTester
+from infra import ComparisonConfig, ModelTester, RunMode
 
 
 class MNISTCNNTester(ModelTester):
     """Tester for MNIST CNN model."""
 
-    def __init__(self, cls):
-        self._model_class = cls
-        super().__init__()
+    def __init__(
+        self,
+        model_class,
+        comparison_config: ComparisonConfig = ComparisonConfig(),
+        run_mode: RunMode = RunMode.INFERENCE,
+    ) -> None:
+        self._model_class = model_class
+        super().__init__(comparison_config, run_mode)
path: tests/jax/models/mnist/cnn/tester.py
user: login: kmitrovicTT
body: Wow, how did we miss this?
created_at: 2025-03-17T10:37:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -7,15 +7,20 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from infra import ModelTester
+from infra import ComparisonConfig, ModelTester, RunMode
 
 
 class MNISTCNNTester(ModelTester):
     """Tester for MNIST CNN model."""
 
-    def __init__(self, cls):
-        self._model_class = cls
-        super().__init__()
+    def __init__(
+        self,
+        model_class,
+        comparison_config: ComparisonConfig = ComparisonConfig(),
+        run_mode: RunMode = RunMode.INFERENCE,
+    ) -> None:
+        self._model_class = model_class
+        super().__init__(comparison_config, run_mode)
path: tests/jax/models/mnist/cnn/tester.py
user: login: mrakitaTT
body: I guess because the training tests are skipped so that code path never run
created_at: 2025-03-17T10:42:05Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-16T23:58:02Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-XLA Tests</td><td>619 ran</td><td>448 passed</td><td>171 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-XLA Tests"] by mikepenz/action-junit-report -->
author_association: CONTRIBUTOR
user: login: codecov-commenter
created_at: 2025-03-16T23:58:23Z
body: ## [Codecov](https://app.codecov.io/gh/tenstorrent/tt-xla/pull/361?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 77.97%. Comparing base [(`cd8b1e8`)](https://app.codecov.io/gh/tenstorrent/tt-xla/commit/cd8b1e83f8f9d4373a0707cbed3fd6288c35ccef?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) to head [(`ba01693`)](https://app.codecov.io/gh/tenstorrent/tt-xla/commit/ba01693c3dacdb8f0759d1770eac5a72af92a251?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

:white_check_mark: All tests successful. No failed tests found.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #361   +/-   ##
=======================================
  Coverage   77.97%   77.97%           
=======================================
  Files          21       21           
  Lines        1108     1108           
=======================================
  Hits          864      864           
  Misses        244      244           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tenstorrent/tt-xla/pull/361?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

author_association: NONE