Merge branch 'main' into ir_mul

Author: vbharadwaj-bk

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## Latest Changes
 
+### v0.6.4 (2026-03-05) 
+Bugfix: added missing MLIR lowerings for
+a pair of JAX primitives (thanks @teddykoker!) 
+
 ### v0.6.3 (2025-02-23) 
 OpenEquivariance v0.6.3 brings long-needed improvements to the
 PyTorch frontend. We strongly encourage all users to upgrade

openequivariance/openequivariance/jax/jvp/conv_prim.py

@@ -161,6 +161,16 @@ def conv_fwd_jvp_abstract_eval(
 
 conv_fwd_jvp_p.def_impl(conv_fwd_jvp_impl)
 conv_fwd_jvp_p.def_abstract_eval(conv_fwd_jvp_abstract_eval)
+mlir.register_lowering(
+    conv_fwd_jvp_p,
+    mlir.lower_fun(conv_fwd_jvp_impl, multiple_results=False),
+    platform="cuda",
+)
+mlir.register_lowering(
+    conv_fwd_jvp_p,
+    mlir.lower_fun(conv_fwd_jvp_impl, multiple_results=False),
+    platform="rocm",
+)
 
 
 # ==============================================================================
@@ -285,6 +295,16 @@ def conv_bwd_jvp_abstract_eval(
 
 conv_bwd_jvp_p.def_impl(conv_bwd_jvp_impl)
 conv_bwd_jvp_p.def_abstract_eval(conv_bwd_jvp_abstract_eval)
+mlir.register_lowering(
+    conv_bwd_jvp_p,
+    mlir.lower_fun(conv_bwd_jvp_impl, multiple_results=True),
+    platform="cuda",
+)
+mlir.register_lowering(
+    conv_bwd_jvp_p,
+    mlir.lower_fun(conv_bwd_jvp_impl, multiple_results=True),
+    platform="rocm",
+)
 
 
 # ==============================================================================

openequivariance/openequivariance/jax/jvp/tp_prim.py

@@ -132,6 +132,16 @@ def tp_fwd_jvp_abstract_eval(X, Y, W, dX, dY, dW, *, L3_dim, kernel, hash):
 
 tp_fwd_jvp_p.def_impl(tp_fwd_jvp_impl)
 tp_fwd_jvp_p.def_abstract_eval(tp_fwd_jvp_abstract_eval)
+mlir.register_lowering(
+    tp_fwd_jvp_p,
+    mlir.lower_fun(tp_fwd_jvp_impl, multiple_results=False),
+    platform="cuda",
+)
+mlir.register_lowering(
+    tp_fwd_jvp_p,
+    mlir.lower_fun(tp_fwd_jvp_impl, multiple_results=False),
+    platform="rocm",
+)
 
 
 # ==============================================================================
@@ -225,7 +235,16 @@ def tp_bwd_jvp_abstract_eval(X, Y, W, dZ, tX, tY, tW, tdZ, *, kernel, hash):
 
 tp_bwd_jvp_p.def_impl(tp_bwd_jvp_impl)
 tp_bwd_jvp_p.def_abstract_eval(tp_bwd_jvp_abstract_eval)
-
+mlir.register_lowering(
+    tp_bwd_jvp_p,
+    mlir.lower_fun(tp_bwd_jvp_impl, multiple_results=True),
+    platform="cuda",
+)
+mlir.register_lowering(
+    tp_bwd_jvp_p,
+    mlir.lower_fun(tp_bwd_jvp_impl, multiple_results=True),
+    platform="rocm",
+)
 
 # ==============================================================================
 # 9. Transpose Rule for Backward JVP

openequivariance/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "openequivariance"
-version = "0.6.3"
+version = "0.6.4"
 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.3"
+version = "0.6.4"
 
 authors = [
   { name="Austin Glover" },