Modified changelog before release.

Author: vbharadwaj-bk

CHANGELOG.md

@@ -1,5 +1,12 @@
 ## Latest Changes
 
+### v0.6.5 (2026-03-22) 
+This release brings `ir_mul` layout support for
+OpenEquivariance. Pass the parameter
+`layout='ir_mul'` to any `TPProblem` instance to use
+a transposed layout for the input and output
+irreps. 
+
 ### v0.6.4 (2026-03-05) 
 Bugfix: added missing MLIR lowerings for
 a pair of JAX primitives (thanks @teddykoker!) 

openequivariance/openequivariance/benchmark/benchmark_utils.py

@@ -114,9 +114,7 @@ def benchmark_forward(
     flops = flops_forward(problem, batch_size=batch_size)
 
     # DATA
-    memory_streamed = memory_streamed_forward(
-        problem, batch_size=batch_size
-    )
+    memory_streamed = memory_streamed_forward(problem, batch_size=batch_size)
 
     result |= calculate_performance_statistics(
         problem=problem,
@@ -172,9 +170,7 @@ def benchmark_backward(
 
     flops = flops_backward(tpp=problem, batch_size=batch_size)
 
-    memory_streamed = memory_streamed_backward(
-        tpp=problem, batch_size=batch_size
-    )
+    memory_streamed = memory_streamed_backward(tpp=problem, batch_size=batch_size)
 
     result |= calculate_performance_statistics(
         problem=problem,
@@ -231,9 +227,7 @@ def benchmark_double_backward(
 
     flops = flops_backward(tpp=problem, batch_size=batch_size)
 
-    memory_streamed = memory_streamed_backward(
-        tpp=problem, batch_size=batch_size
-    )
+    memory_streamed = memory_streamed_backward(tpp=problem, batch_size=batch_size)
 
     result |= calculate_performance_statistics(
         problem=problem,

openequivariance/openequivariance/benchmark/metrics.py

@@ -9,9 +9,7 @@
 logger = getLogger()
 
 
-def memory_streamed_forward(
-    tpp: TPProblem, batch_size: int
-) -> dict[str, int]:
+def memory_streamed_forward(tpp: TPProblem, batch_size: int) -> dict[str, int]:
     """
     This represents an absolute minimum amount of memory that could be streamed on an ideal machine
     It returns the number of bytes streamed total and from each source
@@ -69,9 +67,7 @@ def flops_forward(tpp: TPProblem, batch_size: int) -> dict:
         )
 
     flops_count["CG_decomposition"] *= 3 * batch_size
-    flops_count["linear_combination"] *= (
-        batch_size  # Weights do not require FMA here
-    )
+    flops_count["linear_combination"] *= batch_size  # Weights do not require FMA here
 
     flops_count["total"] = sum(flops_count.values())
     return flops_count

openequivariance/openequivariance/benchmark/plotting/plot_layout.py

@@ -49,7 +49,9 @@ def _dtype_sort_key(dtype_name: str) -> int:
 
     dtype_order = sorted(dtype_order, key=_dtype_sort_key)
 
-    directions = [d for d in ["forward", "backward"] if any(d in grouped[x] for x in grouped)]
+    directions = [
+        d for d in ["forward", "backward"] if any(d in grouped[x] for x in grouped)
+    ]
     if not directions:
         raise ValueError("No forward/backward layout benchmark entries found to plot.")
 
@@ -98,7 +100,9 @@ def _dtype_sort_key(dtype_name: str) -> int:
     fig.supxlabel("Problem")
 
     handles, labels = axs[0][0].get_legend_handles_labels()
-    unique = [(h, l) for i, (h, l) in enumerate(zip(handles, labels)) if l not in labels[:i]]
+    unique = [
+        (h, l) for i, (h, l) in enumerate(zip(handles, labels)) if l not in labels[:i]
+    ]
     if unique:
         axs[0][0].legend(*zip(*unique))
 
@@ -114,6 +118,11 @@ def _dtype_sort_key(dtype_name: str) -> int:
                 if values["mul_ir"] > 0:
                     ratios.append(values["ir_mul"] / values["mul_ir"])
             if ratios:
-                stats = [np.min(ratios), np.mean(ratios), np.median(ratios), np.max(ratios)]
+                stats = [
+                    np.min(ratios),
+                    np.mean(ratios),
+                    np.median(ratios),
+                    np.max(ratios),
+                ]
                 stats_fmt = [f"{val:.3f}" for val in stats]
                 print("\t".join([dtype, direction] + stats_fmt))

openequivariance/openequivariance/core/ConvolutionBase.py

@@ -1,5 +1,3 @@
-import copy
-
 import numpy as np
 
 from openequivariance.core.logging import bcolors, getLogger

openequivariance/openequivariance/core/LoopUnrollTP.py

@@ -1,7 +1,5 @@
 import json
 
-import numpy as np
-
 from openequivariance.core.logging import getLogger
 from openequivariance.core.ComputationSchedule import (
     ComputationSchedule,

openequivariance/openequivariance/core/TensorProductBase.py

@@ -180,4 +180,3 @@ def benchmark_double_backward(
             mode=mode,
             kernel_names=kernel_names,
         )
-

openequivariance/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "openequivariance"
-version = "0.6.4"
+version = "0.6.5"
 authors = [
   { name="Austin Glover" },
   { name="Vivek Bharadwaj" },

openequivariance_extjax/pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "openequivariance_extjax"
-version = "0.6.4"
+version = "0.6.5"
 
 authors = [
   { name="Austin Glover" },

tests/input_validation_test.py

@@ -140,16 +140,14 @@ def test_cpp_checks_forward_dtype(executable_and_buffers, subtests):
                 executable(*buffers)
 
 
-def test_ir_mul_rejects_uvw_problem(dtype):
+def test_ir_mul_rejects_uvw_problem():
     problem = TPProblem(
         "5x5e",
         "1x3e",
         "5x5e",
         [(0, 0, 0, "uvw", True)],
         shared_weights=False,
         internal_weights=False,
-        irrep_dtype=dtype,
-        weight_dtype=dtype,
         layout="ir_mul",
     )