From 38540b48654b9b708befab800edd58f1d82567d2 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:26:09 -0700 Subject: [PATCH 1/4] amd syncwarp fix --- openequivariance/openequivariance/templates/jinja_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/openequivariance/openequivariance/templates/jinja_utils.py b/openequivariance/openequivariance/templates/jinja_utils.py index 609ba3c1..b4ede597 100644 --- a/openequivariance/openequivariance/templates/jinja_utils.py +++ b/openequivariance/openequivariance/templates/jinja_utils.py @@ -26,7 +26,11 @@ def get_jinja_environment(is_hip=False): env.globals["enumerate"] = enumerate env.globals["is_hip"] = is_hip - env.globals["syncwarp"] = "__threadfence_block()" if is_hip else "__syncwarp()" + env.globals["syncwarp"] = ( + '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");__builtin_amdgcn_wave_barrier();__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");' + if is_hip + else "__syncwarp()" + ) env.globals["atomic_add"] = "unsafeAtomicAdd" if is_hip else "atomicAdd" if is_hip: From 723182d043735850978b5215e1c84cca71aad7b0 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 6 Sep 2026 20:22:14 -0700 Subject: [PATCH 2/4] test cases --- .../openequivariance/benchmark/problems.py | 66 +++++++++++++++++++ tests/batch_test.py | 2 + tests/conv_test.py | 6 +- 3 files changed, 73 insertions(+), 1 deletion(-) diff --git a/openequivariance/openequivariance/benchmark/problems.py b/openequivariance/openequivariance/benchmark/problems.py index b486941c..ff359157 100644 --- a/openequivariance/openequivariance/benchmark/problems.py +++ b/openequivariance/openequivariance/benchmark/problems.py @@ -196,6 +196,61 @@ def __init__( ) +class NequIPTPP(TPProblem): + """ + Taken from nequip.nn.interaction_block.InteractionBlock: + https://github.com/mir-group/nequip/blob/27d9d2182da918ab7be0017d8300e53278f5e00e/nequip/nn/interaction_block.py#L89-L116 + + Produces the same set of instructions as ChannelwiseTPP, but in a different + order: the output irreps are sorted, while the instruction list stays in + in1-major / in2-minor enumeration order (ChannelwiseTPP re-sorts it by + output index) + """ + + def __init__( + self, + feature_irreps_in: Irreps, + irreps_edge_attr: Irreps, + feature_irreps_out: Irreps, + label: Optional[str] = None, + irrep_dtype=np.float32, + weight_dtype=np.float32, + ): + feature_irreps_in = Irreps(feature_irreps_in) + irreps_edge_attr = Irreps(irreps_edge_attr) + feature_irreps_out = Irreps(feature_irreps_out) + + irreps_mid = [] + instructions = [] + for i, (mul, ir_in) in enumerate(feature_irreps_in): + for j, (_, ir_edge) in enumerate(irreps_edge_attr): + for ir_out in ir_in * ir_edge: + if ir_out in feature_irreps_out: + k = len(irreps_mid) + irreps_mid.append((mul, ir_out)) + instructions.append((i, j, k, "uvu", True)) + + irreps_mid = Irreps(irreps_mid) + irreps_mid, p, _ = irreps_mid.sort() + + instructions = [ + (i_in1, i_in2, p[i_out], mode, train) + for i_in1, i_in2, i_out, mode, train in instructions + ] + + super().__init__( + feature_irreps_in, + irreps_edge_attr, + irreps_mid, + instructions, + internal_weights=False, + shared_weights=False, + label=label, + irrep_dtype=irrep_dtype, + weight_dtype=weight_dtype, + ) + + FCTPP = FullyConnectedTPProblem CTPP = ChannelwiseTPP @@ -347,6 +402,17 @@ def nequip_problems(): ] +# https://github.com/mir-group/nequip/blob/27d9d2182da918ab7be0017d8300e53278f5e00e/nequip/model/nequip_models.py#L30-L58 +def nequip_oam_problems(): + sh = "1x0e+1x1o+1x2e+1x3o" + hidden = "128x0e+64x1o+32x2e+32x3o" + return [ + NequIPTPP("32x0e", sh, hidden, "nequip-oam-l-first-layer"), + NequIPTPP(hidden, sh, hidden, "nequip-oam-l-main-layers"), + NequIPTPP(hidden, sh, "128x0e", "nequip-oam-l-last-layer"), + ] + + # https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml def nequix_problems(): return [ diff --git a/tests/batch_test.py b/tests/batch_test.py index 7ec6333b..e5af61b3 100644 --- a/tests/batch_test.py +++ b/tests/batch_test.py @@ -14,6 +14,7 @@ diffdock_problems, e3nn_torch_tetris_poly_problems, mace_problems, + nequip_oam_problems, nequip_problems, ) from pytest_check import check @@ -131,6 +132,7 @@ class TestProductionModels(TPCorrectness): production_model_tpps = ( mace_problems() + nequip_problems() + + nequip_oam_problems() + e3nn_torch_tetris_poly_problems() + diffdock_problems() ) diff --git a/tests/conv_test.py b/tests/conv_test.py index 446d0f3f..9ee2c8f2 100644 --- a/tests/conv_test.py +++ b/tests/conv_test.py @@ -19,6 +19,7 @@ mace_problems, diffdock_problems, e3tools_problems, + nequip_oam_problems, ) @@ -172,7 +173,10 @@ def test_tp_triple_bwd(self, conv_object, graph, with_jax): class TestProductionModels(ConvCorrectness): production_model_tpps = ( - mace_problems() + diffdock_problems() + [e3tools_problems()[0]] + mace_problems() + + diffdock_problems() + + [e3tools_problems()[0]] + + nequip_oam_problems() ) @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class") From 26c491360e1886c4812f6887d9ddc26febb90d99 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:49:35 -0700 Subject: [PATCH 3/4] mace-torch pins an old version of e3nn causing issues --- openequivariance/pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 1ea7c89b..089ebdf4 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -56,7 +56,6 @@ dev = [ "pytest", "pytest-check", "pytest-subtests", - "mace-torch", "torch_geometric", "cmake" ] From 7704c1a2df83c1b82b9ad257249486239139114c Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:50:47 -0700 Subject: [PATCH 4/4] shorted comment to just attribute source --- openequivariance/openequivariance/benchmark/problems.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/openequivariance/openequivariance/benchmark/problems.py b/openequivariance/openequivariance/benchmark/problems.py index ff359157..1dcdb173 100644 --- a/openequivariance/openequivariance/benchmark/problems.py +++ b/openequivariance/openequivariance/benchmark/problems.py @@ -200,11 +200,6 @@ class NequIPTPP(TPProblem): """ Taken from nequip.nn.interaction_block.InteractionBlock: https://github.com/mir-group/nequip/blob/27d9d2182da918ab7be0017d8300e53278f5e00e/nequip/nn/interaction_block.py#L89-L116 - - Produces the same set of instructions as ChannelwiseTPP, but in a different - order: the output irreps are sorted, while the instruction list stays in - in1-major / in2-minor enumeration order (ChannelwiseTPP re-sorts it by - output index) """ def __init__(