OpenEquivariance forces are not the gradient of the model energy for l_max = 3 NequIP models

[iansteeg]

Summary

Applying the enable_OpenEquivariance model modifier to NequIP-OAM-L-0.1
leaves the predicted energy unchanged (5e-9 relative, per frame) but changes
the predicted forces by ~1-4 % RMS. Only one of the two force fields can be
the gradient of the shared energy. We measured which, by finite differences of
the model's own output, and it is not the OpenEquivariance one:

• the reference (PyTorch/e3nn) path satisfies dF_i/dx_j = dF_j/dx_i to the
  numerical floor of the measurement;
• the OpenEquivariance path has a Jacobian asymmetry that is independent of
  the finite-difference step over a decade of h, i.e. a real violation, not a
  discretization artifact.

A non-symmetric force Jacobian means curl F != 0, so the forces are not
conservative. In molecular dynamics this injects energy: NVE simulations with two
independently trained l_max = 3 models heat steadily, and the heating rate does
not improve as the timestep is reduced. Rebuilding one of those LAMMPS
artifacts from byte-identical weights with the modifier removed stops the heating
entirely — 79.7 eV/ns becomes a residual consistent with zero, while both
kernels still return the same step-0 total energy to all eight printed figures.

S (l_max = 1), M (l_max = 2) and XL (l_max = 4) route through the same
library and are unaffected, agreeing between the two paths to two to three
significant digits on every test. M, L and XL were each measured as released,
with the OpenEquivariance modifier counted live on the model; the l_max = 1
result is from a fine-tuned S rather than the released S. The defect is specific to L, the only
released model whose highest degree is odd (l_max = 3).

We have not identified the faulty code path. The evidence below is meant to
localize it.

---

Environment

openequivariance	0.6.8
nequip	0.18.0
torch	2.9.1+rocm6.3
python	3.12.13
GPU	AMD Instinct MI250X, gfx90a:sramecc+:xnack-
host	LUMI-G (Cray EX), containerized LAMMPS/NequIP image
model dtype	float32 (model_dtype: float32; the packages cannot be run in float64)

OpenEquivariance falls back to JIT compilation on this platform, emitting:

PyTorch version 2.9.1+rocm6.3 is < 2.10, minimum required for precompiled extension.
HIP does not support precompiled extension yet.

We have tested only this configuration. See "What we have not tested".

Version currency

Checked 2026-08-07, so that this is not read as a stale-version report.

OpenEquivariance. 0.6.8 (released 2026-06-14) is the latest release.
main is ahead of v0.6.8 by exactly one commit, 53b9d45 "Fix JAX tensor
product backward buffer initialization (#205)", which touches only
openequivariance_extjax/src/libjax_tp_jit.cpp and a new
tests/tp_adjoint_test.py. No PyTorch, CUDA or HIP source is modified by it, so
nothing on the code path we exercise has changed since the version we ran.

nequip. Our measurements ran in a container carrying 0.18.0; the current
release is 0.19.0 (main is 2 commits ahead of it, both documentation-only:
d5535a19, e07489d4). Between 0.18.0 and current main the only change to
the OpenEquivariance path is two lines in nequip/nn/_tp_scatter_oeq.py:

class OpenEquivarianceTensorProductScatter(TensorProductScatter):
    _nequip_custom_ops_libs = ("openequivariance",)

(commit 04cd07ee, 2026-05-27). It is consumed in exactly one place,
nequip/nn/graph_model.py:142, where it is collected into the model's AOTI
metadata as an import hint. The tensor-product kernel math is untouched. The
one other backward-related change in that window, b023f19e "Implement
stable_silu_backward and update decomposition table", substitutes
aten.silu_backward in the core-aten decomposition table returned by
nequip_decomp_table() and used only by nequip_make_fx — an overflow guard on
an activation's double-backward, on the export/compile path, applied
identically at every capacity. It cannot produce an L-only first-gradient defect
in eager mode.

Scope of that check. The above establishes version-inertness for the
eager path our finite-difference measurement uses, which is where every
gradient number in this report comes from. It does not establish it for
compiled/AOT dispatch: _nequip_custom_ops_libs exists precisely to affect AOTI
loading, and the LAMMPS ML-IAP artifacts under "Practical consequence" go through
that path. We have not repeated the comparison there.

Affected models

Read directly off the released .nequip.zip artifacts:

model	params	r_max	max l in irreps	edge attribute irreps	affected
NequIP-OAM-S-0.1	616,360	4.5	1	1x0e+1x1e	no
NequIP-OAM-M-0.1	3,189,672	6.0	2	1x0e+1x1o+1x2e	no
NequIP-OAM-L-0.1	9,605,688	6.0	3	1x0e+1x1o+1x2e+1x3o	YES
NequIP-OAM-XL-0.1	32,136,296	6.0	4	1x0e+1x1o+1x2e+1x3o+1x4e	no

The M, L and XL rows were measured on the released artifacts themselves. The S
row was measured on a fine-tuned model of the same architecture — same modules,
same degrees, different weights — since we did not run the released S.

L is the only model whose top-degree irrep has odd parity (1x3o). S is
l=1 with even parity (1x1e); M and XL top out at even degree. If the bug
is in a dispatch or kernel-selection path keyed on the highest degree, or on
odd-parity high-l products, that is where we would look first.

Module counts read off the same artifacts: XL has the same number of
TensorProductScatter modules as L (6 each; M has 4, S has 2), yet XL is
clean. So neither depth, module count, nor parameter count is the discriminant —
XL carries 3.3x L's parameters and one degree more — and the defect is not
monotone in capacity, since all three of those would have made XL the worst
case. Something about the top irrep separates them. We cannot say which reading:
"odd top-degree parity" and "literal l_max = 3" agree on all four released
models, so the four capacities do not distinguish them.

After the modifier we count 12 live OpenEquivariance modules on L, 12 on XL,
8 on M and 4 on S — counted on each live model, XL included, so its clean result
is a real negative and not an unmodified model.

The measurement

If F = -grad E, the Jacobian of the force field is a Hessian and therefore
symmetric. We central-difference the analytic forces and report

A(h) = || J - J^T ||_F  /  || J + J^T ||_F

on a block of 16-24 displaced components, scanning the step h.

We test the Jacobian rather than differencing the energy because model_dtype
is float32 and the total energy is ~1e3-1e4 eV, so -dE/dx carries a noise
floor around 3e-3 eV/A. Forces are O(1) eV/A, so differencing them is about four
orders of magnitude better conditioned. We report both; the Jacobian test is the
sharp one.

The step scan is the evidence, not any single number. A path that truly is a
gradient shows a V: A(h) falls as h^2 (truncation) until float32 evaluation
noise takes over and it rises as 1/h. A path with a genuine violation shows a
plateau at the size of the violation.

Displaced components are drawn from a single local cluster (all within one
cutoff of a seed atom). Components spread across the cell are mutually beyond
the cutoff, the Jacobian block is then diagonal, and its symmetry holds
trivially — that version of the test measures nothing.

Result on the public NequIP-OAM-L-0.1

240-atom periodic polyamide-6 melt (H144 C72 N12 O12), 16-24 components:

h [A]	3e-2	1e-2	3e-3	1e-3	3e-4	1e-4	3e-5	1e-5
reference	3.68e-4	4.17e-5	2.93e-5	6.68e-5	2.12e-4	5.40e-4	1.74e-3	4.34e-3
OpenEquiv	5.72e-4	3.90e-4	3.84e-4	3.82e-4	4.35e-4	7.00e-4	1.72e-3	4.51e-3

The reference path traces the textbook V and bottoms at 2.9e-5. The
OpenEquivariance path is flat at 3.8-3.9e-4 across h = 1e-2 to 1e-3 — a
factor 13 higher, and unchanged when h changes by a decade. Note the two curves
coincide on the noise branch (h <= 3e-5), which shows the two paths have
the same evaluation noise floor: the plateau is signal, not extra noise.

Same model, other structures:

structure	reference (min)	OpenEquiv (min)	ratio
polyamide-6 melt, 240 atoms	2.93e-5	3.82e-4	13.0
gibbsite + water + monomer, 256 atoms	2.97e-5	2.68e-4	9.0
liquid water, 288 atoms	2.99e-5	7.94e-5	2.7

Direct -dE/dx check on the same structures (RMS of F_analytic - (-dE/dx),
eV/A), which is noise-limited at ~3e-3 but corroborates:

structure	reference	OpenEquiv
polyamide-6 melt	2.55e-3 (at floor)	3.46e-2
gibbsite + water + monomer	3.09e-3 (at floor)	1.84e-2
liquid water	3.28e-3 (at floor)	3.74e-3 (unresolved)

Controls

Everything below was run with the same script and the same code path, on the same
three structures — and, for the perturbed-component set, on identical (atom,
axis) triples, since the models sharing r_max = 6.0 select the same ones. The
S/M/L measurements and the XL measurement come from two separate jobs on the same
machine, container and package versions.

1. S, M and XL are clean, each by the conservativity measurement itself. Both
   paths agree to 2-3 significant digits at every h. Minimum-over-scan ratios
   (OpenEquiv / reference) on the three structures: S 1.01, 0.97, 1.01 and
   M 1.03, 0.97, 1.03 on fine-tuned models of those architectures; XL 1.01,
   0.99, 1.02 (fine-tuned) and 0.97, 1.03, 1.03 (released
   NequIP-OAM-XL-0.1, unmodified). The released M and L were run as shipped as
   well (M 1.00, 0.94, 1.04; L 9.01, 13.03, 2.65), so within the released set
   alone the ladder M-clean / L-broken / XL-clean holds with no fine-tuning
   anywhere in it. The released S was not put through this test — the
   l_max = 1 evidence is the fine-tuned S cell above plus the 32 S evaluation
   cells below. Fine-tuning changes weights only, not architecture, but we state
   which artifact each number came from. XL stays within [0.86, 1.21] at
   every one of the eight step sizes, and both of its paths trace the same V —
   large-h log-log slopes agreeing to +-0.02, small-h branch at -0.78 to -0.85 —
   bottoming out at 3.1e-5 to 5.2e-5, inside the S/M floor band of 3.05e-5 to
   8.01e-5. L's OpenEquiv path instead runs flat (slope +0.07 to +0.24) and
   bottoms at 3.0e-4 to 3.7e-4. Across 96 S/M/XL evaluation cells the largest
   RMS force difference between the two paths is 0.0019 meV/A — float32
   noise. Across 32 L cells the smallest is 9.3 meV/A.
2. Energies are identical, per frame. Over 156 independent structures and 12
   model variants, max |E_OEQ - E_ref| / |E| = 2.6e-8 to 4.2e-8, median
   ~7e-9. So both paths evaluate one and the same E(x); the disagreement is in
   the backward pass alone.
3. The modifier really is applied. We count OpenEquivariance modules on the
   live model (4 for S, 8 for M, 12 for L, 12 for XL) rather than assuming —
   nequip.model.modify discards the modifier's return value and relies on
   in-place mutation. The count on XL matters most: it is what makes the clean
   XL result a negative rather than an untested model.
4. Not a neighbor-list artifact. Repeating with the neighbor list frozen
   across all displaced geometries reproduces the same picture.
5. Not run-to-run nondeterminism. Repeat evaluation of one geometry agrees to
   3e-6 eV/A in the forces; the plateau is ~300x that.
6. Not a fine-tuning artifact. The numbers above are from the released
   NequIP-OAM-L-0.1 with no modification. We ran the same step scan on one
   fine-tuned L and one from-scratch L derived from it and saw the same
   plateau signature (min-over-scan ratios 10.8 and 7.5 on the gibbsite
   structure). Separately, on a 128-cell evaluation sweep, all 32 L-capacity
   cells show a force difference between the two paths of 9.3-32.8 meV/A, while
   all 96 S/M/XL cells sit at 0.0007-0.0019 meV/A (float32 noise).

Where the discrepancy lives

Per-species RMS of delta = F_OEQ - F_ref on the released FM-L over 156
structures (44,200 atoms):

species	RMS(delta) [meV/A]	RMS(F_ref) [meV/A]	ratio	atoms
N	187.6	2562.6	0.0732	390
C	59.4	2329.2	0.0255	2340
O	17.1	2518.9	0.0068	13884
H	7.8	1690.2	0.0046	25506
Al	7.6	2311.3	0.0033	2080

Nitrogen and carbon are 6.2 % of the atoms and carry 79 % of sum(delta^2);
nitrogen alone carries 50 %. This is a species effect, not a system effect:
N carries essentially the same delta inside the polyamide melt (185.8 meV/A) as
in the other five systems (194.4 meV/A). The reference force magnitudes are
comparable across C/N/O (2.3-2.6 eV/A), so this is not a normalization artifact.

We do not have an explanation for the nitrogen concentration. It may simply
track local geometry (3-coordinate planar amide N) rather than element identity,
i.e. particular angular configurations that excite the high-l paths hardest.

A third, independent symptom

nequip-compile --mode aotinductor --target pair_nequip on one of our
fine-tuned L models fails its own eager-vs-compiled consistency check, on the
virial specifically:

AssertionError: Compilation check MaxAbsError: 0.000151142 (tol: 5e-05)
for field `virial`. ... the largest absolute (MaxAbs) entry of the model
prediction is 68.5081

The virial is a strain derivative, i.e. another output of the backward pass, and
it moves roughly an order of magnitude more than the forces under this defect
(in our evaluation sweep, force MAE changes ~10 % while stress MAE changes
~40 %). We report it as a third symptom of the same backward-pass discrepancy,
not as an independent bug.

Practical consequence, and a direct causal test in MD

The affected models are used for MD through nequip-prepare-lmp-mliap --modifiers enable_OpenEquivariance. That tool does not bake the kernel into
the weights: the ML-IAP artifact holds the packaged model as raw bytes plus a
modifiers list applied at load time inside LAMMPS. So the same artifact can be
rebuilt with an empty modifiers list and nothing else changed. We did that, and
ran the counterfactual.

Setup, identical for every run below: 6144-atom liquid water (2048 H2O,
40x40x40 A) at 400 K, the same seed configuration, the same velocity seed, the
same LAMMPS deck, NPT then NVT equilibration, then 50 ps of NVE production at
dt = 0.5 fs, with the energy injection fitted over the production stage.

The counterfactual: the same weights, differentiated two ways

The two artifacts differ only in the modifiers list. The packaged model
bytes are byte-identical (equal sha256, equal length), the seed configuration
file is md5-identical, and both kernels return the same step-0 total energy to
all eight figures LAMMPS prints, -32 033.7220 eV — an independent, in-LAMMPS
confirmation of the static finding that the two paths agree on the energy and
differ only in its gradient.

Energy injection over the production stage, three estimators on the same window:

estimator	OpenEquivariance (as deployed)	reference (modifier removed)
six-figure thermo fit, 10 % skip	+79.66 +/- 0.31	-1.86 +/- 0.11
eight-figure log fit, 10 % skip	+79.63 +/- 0.31	-0.29 +/- 0.09
5 x 10 ps segments, mean +/- SEM	+85.1 +/- 16.4	-0.19 +/- 0.92

eV/ns throughout. The three disagree on the reference side for a mundane reason:
the six-figure channel quantizes a ~32 000 eV total at 0.1 eV, and that run takes
only two distinct printed values over the entire 50 ps, so its slope there is
a print-resolution bound rather than a measurement. The eight-figure log channel
is 100x finer, and where both have resolution they agree to 0.1 % (79.658 vs
79.627 on the baseline). The segment estimator is the one whose error bar
accounts for autocorrelation, and it is the one to read.

With the kernel disabled the heating stops. The residual is consistent with
zero and of the opposite sign. Its five segment slopes straddle zero (-0.37,
-0.27, +1.91, -3.45, +1.23) where the OpenEquivariance run's never approach it
(+103.1, +57.1, +86.3, +43.5, +135.4). Accumulated total-energy change across the
production window (5 ps block means, 5 -> 45 ps): -0.010 eV, against
+3.841 eV with the kernel. Peak-to-peak excursion over the whole NVE stage
0.213 eV (sd 0.035) against 4.692 eV (sd 1.176).

State-match controls: the two thermostatted NVT stages agree to 0.037 K, and
the production box volumes to 0.056 %.

One potential-energy surface, differentiated two ways, and only the accelerated
differentiation injects energy. That is the dynamical counterpart of the static
finite-difference result — the forward agrees exactly, the gradient does not, and
the trajectories diverge accordingly.

The second model, and the timestep

A second, independently trained l_max = 3 model — a fine-tune of the released
NequIP-OAM-L; the two are the pair named in control 6 above — run through the
same kernel from the same seed configuration and velocity seed, heats
+232.0 +/- 0.5 eV/ns, about 3x harder. It is also the model with the larger
Hessian asymmetry in the static test above, so the static and dynamic orderings
agree. (It is a different potential and equilibrates 1.7 % denser under NPT, so
read that ratio as ~3x, not to three digits.) It has no reference-path
counterfactual of its own: its heating is attributed to the kernel by analogy
with the case above, not by its own control.

With the kernel enabled, reducing the timestep does not help: 79.7 / 95.4 /
103.3 eV/ns at dt = 0.5 / 0.25 / 0.125 fs, fitted exponent p = -0.19, where
Verlet error on conservative forces would fall as dt^2 (79.7 / 19.9 / 5.0).
That was our first indication the energy came from the force field rather than
the integrator; the counterfactual now shows it directly. There is no
reference-path timestep scan.

Two limits that travel with the causal result

1. It does not show the reference path passes an energy-conservation
   acceptance check. The segment scatter bounds the residual at
   |dE/dt| < 2.8 eV/ns at 95 %, a factor ~1e3 above the 1e-8 to 1e-7 relative
   drift per ns (~3e-4 to 3e-3 eV/ns on this system) that our l_max = 1 and
   l_max = 2 models sustain. A 50 ps window cannot resolve that band.
   "Consistent with zero" is the strongest honest statement; resolving it would
   need a nanosecond-scale reference-path NVE, which we have not run.
2. Ignore the temperature channel. The reference run is a useful control on
   this point: with a provably zero energy injection it still fits a temperature
   slope of +26 to +42 K/ns depending on window, indistinguishable from the
   OpenEquivariance run's. At 50 ps on this system a temperature slope carries at
   least +/-30-40 K/ns, so we quote energy rates only.

Turning the kernel off is not free

At L capacity on this hardware the reference kernel runs this system at
1/3.55 the speed of OpenEquivariance. Both sides are LAMMPS' own whole-run
Performance lines over the identical 200 000-step deck: 3.651 vs 12.970
timesteps/s on the NVE production stage (ratio 3.552; aggregate 3.555). In
per-atom-step terms, 22.43 katom-step/s and 0.158 ns/day on the reference
path against 79.69 katom-step/s and 0.560 ns/day with OpenEquivariance, on 8
GCDs (1 LUMI-G node, 4x MI250X). That is the cost of the only workaround we have,
which is why we would rather see the kernel fixed than route around it.

Related upstream reports

We searched both trackers before writing this (2026-08-07).

mir-group/nequip #586, "[BUG] Strange finetuning training with OEQ +
compile" — open, filed 2026-02-05, no resolution. It reports erratic, unstable
training when fine-tuning NequIP-OAM-L — the same model family and the same
capacity as the defect described here — with OpenEquivariance and compile
enabled together, while either alone is described as "consistent". No
gradient-correctness check is reported there. We raise it because "consistent" is
not "correct": our finite-difference measurement shows the OEQ path is already
non-conservative at L capacity with OEQ alone, and a force field that is not
the gradient of anything is a plausible source of erratic optimization. We offer
this as a hypothesis worth your checking, not as an established identity — we
have not reproduced #586, and that reporter's stack differs from ours in
platform (CUDA), nequip version (0.16.3) and the use of compile.

mir-group/nequip #541, "Add note about OpenEquivariance optimisation
corruption to avoid freezing" — closed — is not this bug. We name it only
to forestall the misidentification: that was a stale JIT cache producing a silent
freeze at model build, fixed by rm -rf ~/.cache/torch_extensions/*. Our runs
build, run, and produce plausible output; the failure mode here is silently wrong
gradients, not a hang, and clearing the cache does not touch it.

Otherwise this appears unreported. In the OpenEquivariance tracker we read
the issue list (38 items) and searched title and body for "conservative",
"non-conservative" and "gradient"; the only gradient-correctness hits are #205
and #138, both buffer-initialization bugs (see below), and neither mentions
forces, energy conservation or MD. In the nequip tracker, "non-conservative",
"not conservative" and "energy drift" return zero results, and of the ten
items mentioning OpenEquivariance only #586 concerns gradients at all. We found
no report in either project of non-conservative forces, kernel-induced energy
drift, or a degree- or parity-specific gradient defect.

A mechanistic lead

We have not read your kernels, so this is reasoning from the symptom and is
offered as somewhere to look rather than as a diagnosis.

Two defects of the same class have already been found and fixed in this codebase:

• PR Zero-initialized L1, L2 gradient buffers to avoid garbage accumulation in gradients #138 (merged 2025-06-21) — "Zero-initialized L1, L2 gradient buffers to
  avoid garbage accumulation in gradients", in the PyTorch/LibTorch path
  (extension/libtorch_tp_jit.cpp, implementations/TensorProduct.py,
  implementations/convolution/TensorProductConv.py): torch::empty /
  torch.empty_like replaced by their zeroed counterparts in the backward.
• PR Fix JAX tensor product backward buffer initialization #205 (merged 2026-07-11, on main, unreleased) — the same defect in the
  JAX FFI path: backward and double-backward kernels accumulate into result
  buffers that are not guaranteed zero.

In both, the forward is untouched and only the gradient is corrupted. That is
precisely our signature: energies from the two paths agree to 5e-9 relative per
frame over 156 structures, forces do not. A buffer or accumulation path reached
only when the instruction list contains a degree-3 (or odd-parity high-degree)
output would explain everything we see at once — forward agreement, gradient
disagreement, and the S/M/XL-clean, L-broken pattern.

We are aware that #138 already zeroed the buffers it names on the PyTorch side,
so this is not a claim that #138 was incomplete. The point is narrower: the class
of defect has now recurred once per backend, #205 shipped with a new adjoint
regression test but that test is JAX-only (tests/tp_adjoint_test.py), and we
are not aware of an equivalent adjoint or gradient-symmetry audit covering the
PyTorch path, and specifically not the HIP JIT-fallback code generation our
runs take. If such a test exists and we have missed it, we would like to be
pointed at it.

Reproducer

oeq_minimal_repro.py, inlined below, is self-contained: it builds its own
periodic NH3/H2O box with ASE and needs only a NequIP package file (e.g. the
public NequIP-OAM-L-0.1.nequip.zip).

oeq_minimal_repro.py (click to expand)

#!/usr/bin/env python
"""Minimal reproducer: OpenEquivariance forces are not the gradient of E.

Self-contained -- builds its own structure with ASE, needs only a NequIP
package file (e.g. the public NequIP-OAM-L-0.1). No project data.

    python oeq_minimal_repro.py NequIP-OAM-L-0.1.nequip.zip
    python oeq_minimal_repro.py NequIP-OAM-M-0.1.nequip.zip   # control

What it measures
----------------
If F = -grad E then the Jacobian of the force field is symmetric,
dF_i/dx_j = dF_j/dx_i (it is a Hessian). We central-difference the analytic
forces and report ||J - J^T|| / ||J + J^T||, with and without
``enable_OpenEquivariance``, scanning the displacement h.

Reading the result
------------------
The scan, not any single number, is the evidence. A path that is a true
gradient shows a V: the asymmetry falls as h^2 (truncation) until float32
evaluation noise takes over and it rises as 1/h. A path whose asymmetry is
*flat* across a decade of h has a real, h-independent violation.

Measured here (openequivariance 0.6.8, torch 2.9.1+rocm6.3, MI250X), min over
the scan:

    NequIP-OAM-L-0.1   reference 1.84e-05   OEQ 7.23e-05   ratio 3.93
    NequIP-OAM-M-0.1   reference 2.78e-05   OEQ 2.75e-05   ratio 0.99
    NequIP-OAM-S-0.1   reference 2.40e-05   OEQ 2.42e-05   ratio 1.01

On real condensed-phase structures the L contrast is larger (up to 13x, with the
OEQ curve flat at 3.8e-4 across a decade of h); this small synthetic box
understates it but discriminates unambiguously.

Displaced components are drawn from ONE local cluster on purpose: components
spread over the cell are mutually beyond the cutoff, the Jacobian block is then
diagonal, and its symmetry is satisfied trivially.
"""
from __future__ import annotations

import sys

import numpy as np
import torch


def build_structure(seed: int = 0):
    """A small periodic NH3/H2O box (90 atoms, 9 N).

    These exact constants are the ones the numbers in the report were measured
    with -- do not retune them casually. The signal depends on local geometry:
    a denser, more N-rich variant (spacing 3.8, half NH3) was tried and gave a
    *weaker* contrast (1.85x rather than 3.93x).
    """
    from ase import Atoms
    from ase.build import molecule

    rng = np.random.default_rng(seed)
    box, spacing = 12.0, 4.0
    atoms = Atoms(cell=[box, box, box], pbc=True)
    n = int(box // spacing)
    for i in range(n):
        for j in range(n):
            for k in range(n):
                m = molecule("NH3" if (i + j + k) % 3 == 0 else "H2O")
                # random orientation, so no accidental symmetry
                m.rotate(rng.uniform(0, 360), rng.normal(size=3), center="COM")
                m.translate(np.array([i, j, k]) * spacing + spacing / 2
                            + rng.normal(scale=0.15, size=3))
                atoms += m
    atoms.wrap()
    return atoms


def build_model(package: str, oeq: bool, device: str):
    from nequip.model import ModelFromPackage, modify
    from nequip.utils.global_state import set_global_state

    set_global_state()
    obj = ModelFromPackage(package, compile_mode="eager").to(device)
    if oeq:
        obj = modify(model=obj,
                     modifiers=[{"modifier": "enable_OpenEquivariance"}])
    # nequip >= 0.18 hands back a ModuleDict keyed `sole_model`
    model = (next(iter(obj.values()))
             if isinstance(obj, (dict, torch.nn.ModuleDict)) else obj)
    model = model.to(device).eval()
    n_oeq = sum("OpenEquivariance" in type(m).__name__
                or "openequivariance" in (type(m).__module__ or "").lower()
                for m in model.modules())
    # `nequip.model.modify` drops the modifier's return value and relies on
    # in-place mutation -- check rather than assume
    assert n_oeq > 0 if oeq else n_oeq == 0, f"OEQ={oeq} but {n_oeq} modules"
    return model, n_oeq


def forces_fn(model, device):
    from nequip.data import AtomicDataDict, from_ase
    from nequip.data.transforms import (ChemicalSpeciesToAtomTypeMapper,
                                        NeighborListTransform)
    md = dict(getattr(model, "metadata", {}) or {})
    mapper = ChemicalSpeciesToAtomTypeMapper(
        model_type_names=str(md["type_names"]).split())
    nl = NeighborListTransform(r_max=float(md["r_max"]))

    def f(atoms):
        d = from_ase(atoms)
        for k in ("total_energy", "forces", "stress", "virial", "free_energy"):
            d.pop(k, None)
        d = nl(mapper(d))
        d = AtomicDataDict.to_(d, device)
        out = model(AtomicDataDict.batched_from_list([d]))
        return out[AtomicDataDict.FORCE_KEY].detach().double().cpu().numpy()
    return f


def asymmetry(f, atoms, comps, h):
    """||J - J^T||_F / ||J + J^T||_F on the chosen component block."""
    n = len(comps)
    J = np.empty((n, n))
    for q, (i, a) in enumerate(comps):
        p, m = atoms.copy(), atoms.copy()
        p.positions[i, a] += h
        m.positions[i, a] -= h
        col = -(f(p) - f(m)) / (2 * h)
        J[:, q] = [col[j, b] for j, b in comps]
    return (np.linalg.norm(J - J.T) / max(np.linalg.norm(J + J.T), 1e-300))


def main() -> None:
    package = sys.argv[1]
    device = "cuda" if torch.cuda.is_available() else "cpu"
    atoms = build_structure()
    print(f"structure: {len(atoms)} atoms, "
          f"{dict(zip(*np.unique(atoms.get_chemical_symbols(), return_counts=True)))}")
    print(f"package  : {package}")
    print(f"device   : {device}\n")

    rng = np.random.default_rng(1)
    i0 = int(rng.integers(len(atoms)))
    near = np.where(atoms.get_distances(i0, np.arange(len(atoms)),
                                        mic=True) <= 6.0)[0]
    flat = rng.choice(np.concatenate([near * 3 + a for a in (0, 1, 2)]),
                      size=min(16, 3 * len(near)), replace=False)
    comps = [(int(c // 3), int(c % 3)) for c in sorted(flat)]
    print(f"{len(comps)} components from a cluster of {len(near)} atoms "
          f"around atom {i0}\n")

    steps = [3e-2, 1e-2, 3e-3, 1e-3, 3e-4, 1e-4]
    res = {}
    for oeq in (False, True):
        model, n = build_model(package, oeq, device)
        f = forces_fn(model, device)
        res[oeq] = [asymmetry(f, atoms, comps, h) for h in steps]
        print(f"{'OEQ ' if oeq else 'ref '}({n} OEQ modules): "
              + "  ".join(f"{v:.2e}" for v in res[oeq]))
        del model

    print()
    print(f"{'h [A]':>10s} {'reference':>12s} {'OpenEquiv':>12s} {'ratio':>8s}")
    for h, a, b in zip(steps, res[False], res[True]):
        print(f"{h:10.1e} {a:12.3e} {b:12.3e} {b / a:8.2f}")
    print()
    print(f"min over scan: reference {min(res[False]):.3e}, "
          f"OEQ {min(res[True]):.3e}, ratio {min(res[True]) / min(res[False]):.2f}")


if __name__ == "__main__":
    main()

python oeq_minimal_repro.py NequIP-OAM-L-0.1.nequip.zip   # affected
python oeq_minimal_repro.py NequIP-OAM-M-0.1.nequip.zip   # control
python oeq_minimal_repro.py NequIP-OAM-S-0.1.nequip.zip   # control

Measured output on this hardware (90-atom box, minimum over the scan):

model	reference	OpenEquiv	ratio
NequIP-OAM-L-0.1	1.84e-5	7.23e-5	3.93
NequIP-OAM-M-0.1	2.78e-5	2.75e-5	0.99
NequIP-OAM-S-0.1	2.40e-5	2.42e-5	1.01

On the small synthetic box the L contrast is weaker than on real
condensed-phase structures (13x, with a flat plateau across a decade of h), but
the discrimination is unambiguous. The signal depends on local geometry: a
denser, more nitrogen-rich variant of the box gave a weaker contrast (1.85x),
so we ship the constants that were measured rather than tuned ones. Our
production numbers above come from real DFT-sampled structures.

What we have not tested

• Any platform other than AMD/HIP (MI250X, gfx90a) with the JIT fallback
  path. We have no CUDA measurement, and the JIT fallback is itself unusual.
• Any OpenEquivariance version other than 0.6.8.
• Any l_max = 3 architecture other than NequIP-OAM-L (we have not isolated
  whether it is l_max = 3 as such, odd top-degree parity, or something
  particular to this model's irrep layout).
• float64. model_dtype is float32 in these packages and casting the model with
  .double() raises expected mat1 and mat2 to have the same dtype, but got: float != double, so we could not raise the precision floor.
• Whether the fault is in the forward tensor-product kernel's backward, in the
  autograd binding around it, or in the fused scatter.
• A reference-path NVE long enough to resolve true energy conservation. Our
  counterfactual bounds the residual at < 2.8 eV/ns over 50 ps; the band we would
  want to confirm is ~1e3 below that and needs nanosecond-scale runs at 3.65
  timesteps/s. Not run.
• A reference-path timestep scan, so we cannot offer "flat with the kernel,
  dt^2 without it" — only "flat with the kernel, and zero without it at
  dt = 0.5 fs".
• The compiled/AOT path, as a version question. Every gradient number here
  comes from eager evaluation under nequip 0.18.0, and "Version currency" above
  establishes that the OEQ path is unchanged between that and current main
  for eager. We did not repeat that check for AOTInductor / torch.compile
  dispatch — which is what the LAMMPS ML-IAP artifacts under "Practical
  consequence" use, and which is exactly what the _nequip_custom_ops_libs
  registration added since 0.18.0 is there to affect. Read the MD numbers as
  measured on our stack, not as established for current main compiled.

What would help us narrow it

If you can point at which contraction paths differ between l_max = 2 and
l_max = 3, or at any place where the backward is derived analytically rather
than by autograd over the forward, or at any backward buffer whose
zero-initialization is conditional on the instruction list, we are happy to run
targeted tests on this hardware and report back. We can also run an adjoint or
gradient-symmetry test of your choosing on MI250X, including at l_max = 3, if
you would rather specify the check than interpret ours.

---

Filed for information also to the NequIP/mir-group maintainers, since the
affected artifacts are released NequIP-OAM foundation models built on top of
OpenEquivariance.

[vbharadwaj-bk]

Thanks for this report! We will triage and take a look.