diff --git a/openequivariance/openequivariance/templates/loop_unroll_tp.cuh b/openequivariance/openequivariance/templates/loop_unroll_tp.cuh
index 52b92191..ea94762e 100644
--- a/openequivariance/openequivariance/templates/loop_unroll_tp.cuh
+++ b/openequivariance/openequivariance/templates/loop_unroll_tp.cuh
@@ -1,4 +1,5 @@
-{%- from 'macros.jinja' import layout_load, layout_store, reg_store with context %}
+{%- from 'macros.jinja' import layout_load, layout_store, reg_store,
+        l2_smem_index with context %}
 {%- from 'wmm.cuh' import generate_matmul %}
 
 {%- macro generate_segment_kernel_forward(id, segment, warp_size) %}
@@ -51,7 +52,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__
 
                 #pragma unroll
                 for(int j = 0; j < {{L2[v].ir.dim}}; j++)
-                    l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}] * weight;
+                    l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}] * weight;
             {%- elif problem.instructions[k].connection_mode == "uvw" %}
                 {# Stream weights here #}
                 {%- set slice_size = L3[w].mul * L1[u].mul %}
@@ -61,7 +62,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__
                 }
                 #pragma unroll
                 for(int j = 0; j < {{L2[v].ir.dim}}; j++)
-                    l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}];
+                    l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}];
             {%- endif %}
 
             // ----------------- CORE CALCULATION -----------------
@@ -184,11 +185,11 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__
             {%- if k == 0 or interactions[k][1] != interactions[k-1][1] or L2[v].mul > 1 or L1[u].mul != L1[interactions[k-1][0]].mul %}
                 #pragma unroll
                 for(int j = 0; j < {{L2[v].ir.dim}}; j++) {
-                    l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}]; 
+                    l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}];
                     l2_grad[j] = 0.0; 
 
                     {%- if double_bwd %}
-                    l2_original[j] = L2_original[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}]; 
+                    l2_original[j] = L2_original[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}];
                     {%- endif %}
                 }
             {%- endif %}
@@ -287,7 +288,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__
                 if(lane_id == 0) {
                     #pragma unroll 
                     for(int j = 0; j < {{L2[v].ir.dim}}; j++)
-                        L2_grad_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}] += l2_grad[j];
+                        L2_grad_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}] += l2_grad[j];
                 }
             {%- endif %}
 
diff --git a/openequivariance/openequivariance/templates/macros.jinja b/openequivariance/openequivariance/templates/macros.jinja
index 59727e88..ce2a840b 100644
--- a/openequivariance/openequivariance/templates/macros.jinja
+++ b/openequivariance/openequivariance/templates/macros.jinja
@@ -66,6 +66,14 @@ Keys map to lists of tuples with (name, dtype, num_elements) of each subarray.
     {%- endif %}
 {%- endmacro %}
 
+{%- macro l2_smem_index(layout, mul, dim, start, mul_var, dim_var) -%}
+    {%- if layout == "ir_mul" -%}
+        {{mul_var}} + {{start}} + {{dim_var}} * {{mul}}
+    {%- else -%}
+        {{dim_var}} + {{start}} + {{mul_var}} * {{dim}}
+    {%- endif -%}
+{%- endmacro %}
+
 {%- macro declare_smem_variables(segment, smem_base) %}
     {%- for name in segment.smem %}
         {%- if name != "total" %}
diff --git a/tests/batch_test.py b/tests/batch_test.py
index 7ec6333b..e5c0bf7b 100644
--- a/tests/batch_test.py
+++ b/tests/batch_test.py
@@ -313,23 +313,22 @@ class TestIrMul(TPCorrectness):
 
     tpps = mace_problems() + [
         oeq.TPProblem(
-            "5x5e",
-            "1x3e",
-            "5x5e",
+            in1,
+            in2,
+            in1,
             [(0, 0, 0, "uvu", True)],
             shared_weights=False,
             internal_weights=False,
-            label="ir_mul_repr_5x1x5_l535",
-        ),
-        oeq.TPProblem(
-            "13x5e",
-            "1x3e",
-            "13x5e",
-            [(0, 0, 0, "uvu", True)],
-            shared_weights=False,
-            internal_weights=False,
-            label="ir_mul_repr_13x1x13_l535",
-        ),
+            label=label,
+        )
+        for in1, in2, label in [
+            ("5x5e", "1x3e", "ir_mul_repr_5x1x5_l535"),
+            ("13x5e", "1x3e", "ir_mul_repr_13x1x13_l535"),
+            ("32x1e", "3x1e", "ir_mul_L2mul3_l111"),
+            ("32x1e", "2x2e", "ir_mul_L2mul2_l121"),
+            ("16x2e", "8x2e", "ir_mul_L2mul8_l222"),
+            ("16x1e", "40x1e", "ir_mul_L2mul40_l111"),
+        ]
     ]
 
     @pytest.fixture(params=tpps, ids=lambda x: x.label, scope="class")
diff --git a/tests/conv_test.py b/tests/conv_test.py
index 446d0f3f..cd8228cf 100644
--- a/tests/conv_test.py
+++ b/tests/conv_test.py
@@ -328,23 +328,22 @@ def conv_object(self, request, problem, extra_conv_constructor_args):
 class TestIrMulLayout(ConvCorrectness):
     production_model_tpps = mace_problems() + [
         oeq.TPProblem(
-            "5x5e",
-            "1x3e",
-            "5x5e",
+            in1,
+            in2,
+            in1,
             [(0, 0, 0, "uvu", True)],
             shared_weights=False,
             internal_weights=False,
-            label="ir_mul_repr_5x1x5_l535",
-        ),
-        oeq.TPProblem(
-            "13x5e",
-            "1x3e",
-            "13x5e",
-            [(0, 0, 0, "uvu", True)],
-            shared_weights=False,
-            internal_weights=False,
-            label="ir_mul_repr_13x1x13_l535",
-        ),
+            label=label,
+        )
+        for in1, in2, label in [
+            ("5x5e", "1x3e", "ir_mul_repr_5x1x5_l535"),
+            ("13x5e", "1x3e", "ir_mul_repr_13x1x13_l535"),
+            ("32x1e", "3x1e", "ir_mul_L2mul3_l111"),
+            ("32x1e", "2x2e", "ir_mul_L2mul2_l121"),
+            ("16x2e", "8x2e", "ir_mul_L2mul8_l222"),
+            ("16x1e", "40x1e", "ir_mul_L2mul40_l111"),
+        ]
     ]
 
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")