Repository navigation
Add the kernel backend contract and the plain-torch reference backend - #1740
Open
aacostadiaz wants to merge 17 commits into
Open
aacostadiaz wants to merge 17 commits into
aacostadiaz wants to merge 17 commits into
Conversation
The first half of CORE-4, and the half that is settled. The Racah expression for `<j1 m1; j2 m2 | j3 m3>` and the Wigner 3j table built on it, in pure numpy, importing neither e3nn nor cuequivariance. The sum is accumulated over Fraction, so the only floating-point work is one square root and one multiplication per coefficient. Against sympy, which is an independent derivation by a different route and a test-only dependency, the agreement over every triple with l <= 5 is 1.1e-16. The second half needs a convention that no ticket states. Turning these coefficients into the real basis the models use needs e3nn's real spherical harmonic ordering, and it is not the standard one. Measured against o3.spherical_harmonics on 3000 random directions, e3nn's l=1 basis is (x, y, z) while the standard m = -1, 0, +1 ordering gives (y, z, x). The relation is a permutation per degree, not a phase, which is why the sign of the 3j table against e3nn's cannot be written as a product of per-degree factors: over the grid up to l=6 it is uniform per triple and takes all four fourth roots of unity, with no factorisation. The same convention is needed by the native spherical harmonics of ARCH-1. Two tickets have to agree on something neither of them writes down, and if they disagree the symptom is numerical rather than an import error. This branch carries only what is verified. The basis itself follows.
The rest of CORE-4. The reduced symmetric tensor-product basis in pure numpy, importing neither e3nn nor cuequivariance, and reproducing the counts the legacy cueq-only path produces. On the measured anchor, irreps_in=0e+1o+2e+3o at correlation 3, the per-order counts are 1, 4, 8 for 0e, 1, 3, 12 for 1o and 1, 5, 14 for 2e. Summed over the body orders a model of that correlation builds, that is 13, 16 and 20, and so 29 for keep_ir=0e+1o. Those are the ticket's numbers, and the 29 is the one a host without cuequivariance silently turned into 86, training a different and more heavily parametrized network for the same hyperparameters. Three things are pinned here because they are the on-disk weight format. The path order is the coupling enumeration, outermost factor first, over the irrep total order defined in irreps.py: ascending degree, even parity before odd. It is a pure function of the declaration, with no dependence on dictionary iteration or on a floating-point comparison. The reduction is a modified Gram-Schmidt in that enumeration order. An SVD would give the same span with an arbitrary basis inside it and a sign that moves between LAPACK builds, which is not a file format. The normalization is unit Frobenius norm per path, sign fixed so the first structurally non-zero entry is positive. The real basis is the textbook one, m = -l..+l, stated in real_basis.py and not inherited. e3nn's is a different orthogonal basis of the same space: it agrees up to a signed permutation at l = 0 and l = 1 and then diverges, mixing m = 0 with m = +2 through a rotation at l = 2. Matching it element for element would mean reimplementing its construction for no reason beyond it being the one that library chose, so v1 states its own and the converter absorbs the difference, as it already has to for the node embedding's factor of sqrt(num_elements). Equivalence is what the oracle tests assert: the spans agree and the change of basis is orthogonal. Nothing observable depends on the choice, because the weights multiplying the basis are learned. tests/architecture/test_no_vendor_deps.py keeps the dependency out by every route that put it there: a static import at any depth, an importlib reach-in, the runtime module table, and the count itself asserted to be 29 whether or not cuequivariance is importable in the environment running the test. Note for later consolidation: irreps.py parses declarations, and so does observables/grammar.py from CORE-1. The target layout puts the algebra here, so the grammar module should delegate once both are on the integration branch.
The rest of CORE-4. The unreduced basis, the full/reduced conversions, the canonical layout boundaries, and the second oracle. The unreduced basis reproduces the legacy numbers too: 28 paths for 0e, 58 for 1o and 73 for 2e on the anchor, so 86 for keep_ir=0e+1o against 29 reduced. Both halves of the gap the CUET_AVAILABLE branch used to open are now derived from first principles. full -> reduced is exact and unique, reduced -> full is exact and under-determined and returns the minimum-norm representative. Neither is bit-exact in weight space and neither can be, because the map is a projection: the full basis carries directions the reduced one does not, and those directions are gauge. So both are tested on values, the round trip agreeing to 1e-14 at fp64, and one test asserts the weights are NOT equal so that nobody later tightens it into something false. One factor was measured rather than assumed. Symmetrizing sums over the permutations instead of averaging, so it carries a factor of correlation!, and the converted weights came out at exactly 3! times the original at correlation 3. Dividing it out is what makes the conversion preserve the function outright rather than up to a scale every caller would have to know about. The layout boundaries cost nothing, as the ticket says: a concatenate, a split, and a reshape. mul_ir to ir_mul refuses a multi-term declaration instead of reshaping it wrongly, because for several terms the two layouts are not related by one reshape and the failure would be silently scrambled weights. The cuequivariance oracle compares spans rather than entries, which follows from the real-basis convention: the two describe the same subspace in different conventions, and the difference is gauge. Same rank, same span, and every path of one solved for as a combination of the other with a residual under 1e-10, on all eight grid points including the anchor. cuequivariance stays a test-only dependency and test_no_vendor_deps.py asserts mace_core never imports it. That last test caught its own first version: it asserted the absence in-process, and the oracle test imports cuequivariance on purpose in the same session. It runs in a fresh interpreter now.
The torch-free half of BKD-1: descriptors, capabilities, the Protocol, the entry-point registry and the canonical weight layout. Both implementations read this and a checkpoint records it, so none of it may import a framework, and the test that says so runs in a fresh interpreter. Three inversions of what the frozen tree does, one per piece. A backend declares, it never chooses. On the frozen tree a host without cuequivariance silently switches the Clebsch-Gordan basis and trains a different, more heavily parametrized network for the same hyperparameters. Here the basis is a recorded field of the descriptor, a backend states which bases it serves, and a mismatch raises with both sides named. A test pins the two counts the field selects, 29 reduced against 86 full, so the coupling cannot come back by another route. Dispatch resolves once. Every method of the Protocol is a factory: a descriptor goes in at model build time, an op comes out, and the op is frozen into the module tree. Nothing in the contract is reachable from forward, which is what removes the six `hasattr(self, "conv_fusion")` branches inside the interaction blocks and the three `.forward` monkeypatches. One canonical layout, mul_ir, with symmetric-contraction weights as a flat [Z, A, mul] array over the reduced basis in the order CORE-4 pins. Conversion happens at the checkpoint boundary and never in forward. The frozen tree needs five conversion command line tools because each backend stores weights its own way. Two smaller decisions worth stating. Discovery records import failures and resolution raises: a machine without a CUDA runtime is expected to carry entry points it cannot load, and listing them has to work exactly there, but asking for one by name cannot fall back to another backend because another backend is another set of numbers. And a backend that declares no double backward is rejected at build time for force or stress training, loudly and never as a warning, since a first derivative taken through a backward that is not itself differentiable gives wrong forces rather than no forces. The torch half follows: the five dispatched ops as custom_op with a differentiable backward, and the plain-torch reference backend.
Three of BKD-1's five dispatched ops, as torch custom operators with a meta implementation and a registered autograd rule: segment_sum, symmetric_contraction and channelwise_tp_conv. torch enters the tree here, in mace-torch and nowhere else, with a floor of 2.4 because that is where torch.library.custom_op and its register_fake and register_autograd arrived. All three pass gradcheck **and** gradgradcheck at fp64, symmetric_contraction at correlation 1, 2 and 3. The second derivative is the requirement that matters: training on forces differentiates a quantity a backward pass produced, and a backward that is not itself differentiable gives wrong forces rather than an error, so a model trained on them looks like it is working. The outer-power backward is the piece that had to be derived rather than inferred. The n-fold outer power differentiates into n terms by the product rule, each contracting the gradient against the power of one degree less on every axis but one, and an error in one term is invisible at correlation 1 and 2. That is why the gradcheck is parametrized over all three. The convolution is always node-level and takes the node count as a plain int. Never edge_index.max(), which is a data-dependent host read: a recompile on every batch and a broken CUDA-graph capture under reduce-overhead. A test pins that a node no edge reaches still appears in the output. Under torch.compile(fullgraph=True, dynamic=True), five batch sizes produce two compiled graphs rather than five. Two rather than one is dynamo's own first-call behaviour, a static graph then the dynamic one it settles on, and not a property of these ops; what would fail the assertion is five. The reference backend registers through an entry point in this package's metadata, so mace_core names no backend and discovery works the same for a third-party one. Not in this commit: linear and fully_connected_tp, the reference backend itself, and the native spherical harmonics and radial bases. The two ops the ticket names for gradcheck are the two here. One note for whoever updates CI. The `packages` job in ci-core.yaml installs the four distributions and says in a comment that it installs no torch, and that this is the point. With this commit mace-torch declares torch, so that comment stops being true. The assertion it was protecting, that mace_core stands up framework-free, is the one tests/architecture/test_no_vendor_deps.py makes directly, and it still holds.
The rest of BKD-1. The equivariant linear op, the native spherical harmonics and radial bases, and the backend wiring every factory. Plain torch, no e3nn. The harmonics are built from this project's own coefficients. Each degree is the previous one tensored with the l = 1 block and projected back with the same wigner_3j_real the symmetric contraction uses. That is deliberate: the harmonics and the contraction have to agree on a convention or the model is wrong, and the symptom is numerical rather than an error. Deriving one from the other makes disagreement impossible instead of making it testable. Measured: unit norm per degree for every direction, and the inner product within a degree invariant under a common rotation to 1.1e-15. The linear map carries its block structure as index tensors rather than as Python constants, which is what lets one registered operator serve every irreps declaration in a model. One weight appears 2l+1 times, once per component of its irrep, and that repetition is the equivariance: a test asserts a 1o input cannot move a 0e output however the weights are set. The full force path differentiates twice. Positions to directions to harmonics to a tensor product reduced onto nodes to a contraction to an energy, with gradcheck and gradgradcheck against the positions. The forces also sum to zero, which the chain gets for free because the energy depends on positions only through differences, and which a mixed-up sender and receiver would break. Canonical weights round trip: a fresh instance loaded from what another wrote reproduces its outputs. For this backend both directions are views, since it holds the canonical layout directly, which is how the format gets defined by something rather than declared in prose. Two places the reference raises instead of guessing. The skip connection builds only against scalar attributes, which is the case the models use, because a wrong skip connection is a wrong model that still trains. And segment reduce is sum only. Both name what is missing. Two of the failures on the way here were in the tests rather than the code, and both are worth the note they now carry. The radial weights have to be as wide as the convolution's coupling path count, read off the built op rather than assumed. And the chain has to be built once: gradcheck calls a function many times, so one that re-randomized its weights per call is a different function each time and fails for a reason that has nothing to do with its derivative. Still open, and not open because of this branch: the parity against the frozen legacy on the tiny anchors needs a v1 model to compare, which is ARCH-2 and ARCH-3, and tests/parity does not exist until PAR-1.
Path order and per-path normalization were stated as the whole of the on-disk weight format. They are not. The enumerated coupling paths are linearly dependent, so which of them survives the reduction is a free choice that neither convention constrains, and two implementations resolving it differently span the same space with vectors no reordering relates. Measured against cuequivariance on the grid BKD-1a uses: four of the five points agree up to a signed permutation, and the fifth, irreps_in 0e+1o+2e+3o to 0e+1o+2e at correlation 3, decomposes into 44 blocks of size one plus one of size two and one of size three, all inside the 2e slot at body order three. Its worst block has condition number 2.79. Every surviving path now carries a CouplingTree: the consumed input slice and the running intermediate irrep at each step, written as "0:0e|1:1o|2:2e" and parsed back. A label names a path wherever it sits, so a reader of a checkpoint places the weight itself instead of trusting that its enumeration matches the writer's. The test that matters rebuilds each basis tensor from its label alone and compares against the stored one; the labels separate by 0.06 in the anchor's worst case, against a 1e-12 tolerance. Also records why the per-path norm is one rather than sqrt(ir.dim): one rule for every output irrep keeps weights on disk at a comparable scale, and the factor a backend or an initializer wants is carried explicitly by that converter. BKD-1a (ACEsuit#1564)
The lint job installs the four packages and ruff, ty and pytest, and nothing else. Both oracle suites guarded a module-level import with `find_spec` and a module-level skip, which is right at runtime and invisible to a type checker: `ty` sees `from sympy.physics.wigner import ...` and `import cuequivariance`, resolves neither, and fails the job on a file that never runs there. `importorskip` says the same thing in one expression and gives the checker a value it can reason about. Same skip, same reason string, and the suites still report a skip rather than an error on a machine without the oracle. Found by running the lint job's own environment against this branch rather than the development one, which has both installed and therefore says nothing.
The basis is numpy arrays end to end and the distribution declared nothing, so `pip install -e packages/mace-core` installed it without numpy. Every job that installs this distribution alone then works or not depending on what else is in the environment: the packages job happened to have numpy pulled in by another package, and the lint job did not, so `ty` reported nine unresolved imports of a module the package cannot run without. Not caught before pushing because the venv it was checked in had numpy left over from installing a different branch's mace-core, which declares it. A check environment that accumulates installs answers a question nobody asked. Rebuilt from scratch, and the failure reproduces on the branch as it stood.
The empty case was keyed on the enumeration finding nothing, and that is not the only way to arrive at no paths. An irrep can be reachable and still be carried entirely by the antisymmetric part: `1o x 1o -> 1e` is the cross product, and a symmetric product has none of it. The paths are enumerated, they cancel under symmetrization, `kept` comes back empty and `np.stack` raises `need at least one array to stack`. A caller asking for such an output irrep has asked for something the symmetric product does not contain, which is an answer and not a failure, so both basis builders return a zero-path array in the shape the caller expects. The failure was reachable through an ordinary loop over several output irreps, where one slot having no paths took the whole request down and the message named neither the irrep nor the reason. Reported by Bugbot on the review of this branch and reproduced before being believed. The unreduced basis was always right here and still is: it has the one path, because it does not symmetrize.
`ty` runs over `packages/` in the lint job and has never been clean on this branch. Seven diagnostics, and one of them is a real defect. `ReferenceFullyConnectedTP.forward` accumulated from `None` and returned it. With no scalar channels to weight by the loop never runs, so the function returns `None` while declaring `Tensor`, and the `None` travels into the rest of the model to fail somewhere with no connection to here. It accumulates from zeros now, which is the same arithmetic when the loop runs and the right answer when it does not: a sum over an empty set is zero. The rest are the checker seeing less than a reader does. `register_buffer` leaves an attribute typed as a `Module`, so reading its shape reads as subscripting a module; the three classes that hold buffers annotate them. A descriptor's read-only property assigned to inside `pytest.raises` goes through `setattr`, since the assignment is the thing under test and a checker is right to reject it written out. And the cueq oracle's module-level import becomes an `importorskip`, which is what the lint job can resolve.
An op allocates its weights at zero and waits to be loaded. That is right for every model that comes from a checkpoint, and it builds a model that cannot be trained: zero-valued tensors multiply to zero, so does the gradient, and a run started from scratch finishes and reports a model that never moved. No test noticed, because every one of them either loads a checkpoint or overwrites the parameters first. So the three ops that own their weights now own the draw as well, beside the canonical pair they already had, and the protocol says so: it named the canonical pair in the ticket and declared neither. The scale is part of the checkpoint format rather than of an implementation. The canonical layout carries the normalization folded into the number, where the frozen tree keeps a standard normal and applies a factor inside the operation, so the factors are stated with the layout and every backend draws at the same size. A linear divides by the square root of the multiplicity feeding its output irrep, the skip connection by the square root of that times the number of element attributes, and the symmetric contraction by nothing. The seed is per op, derived from the run's seed and the op's path. One seed for the whole model gives two ops of the same shape one draw, which is a symmetry that never breaks, and the global generator would make a model depend on how many numbers anything else drew first. The path is hashed rather than passed to `hash`, whose salt changes every process.
…reforge/bkd-1-kernel-backend
9 tasks
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:08 — with 
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:08 — with
GitHub Actions
Failure
GitHub Actions
Failure
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
`tomllib` is 3.11 and up, so a bare import breaks collection on the 3.10 leg of the matrix. pytest declares `tomli` there and this file is collected by pytest or not at all, so the fallback the rest of the tree already uses always resolves.
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:20 — with
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:20 — with
GitHub Actions
Error
GitHub Actions
Error
…'s count Two defects in the reference backend, both found in review of the pull request. The symmetric contraction stacked every output irrep's basis on the path axis. Each one has its own component count on the next axis, so `0e+1o` did not build at all, which is the output every interaction layer asks for. Worse than the failure is the case that does not fail: two irreps of the same size join without complaint and then write into one shared slice instead of adjacent ones, which is a model that trains and is wrong. There is now one contraction per output irrep, concatenated on the component axis, which is the shape the frozen tree reaches by holding one contraction object per irrep. That uncovered a second one underneath it. An output irrep no body order can reach has an empty basis, and numpy cannot infer a trailing extent from an empty array, so a `2e` output failed to build over its own body order one. The extent is computed rather than inferred. The skip connection's weight count was whatever the caller passed, defaulting to zero, while the operation allocated one equivariant linear map per element attribute and never read the field. A capability filter reading it, or a checkpoint sized against it, would have been wrong by the difference: 0 against the 64 the operation holds. It is derived from the irreps now, like every other descriptor's, and the caller-supplied field is gone. The entry-point test reads the installed metadata rather than parsing `pyproject.toml`, which is what the registry actually searches and needs no TOML parser on the 3.10 leg of the matrix.
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:42 — with
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:42 — with
GitHub Actions
Error
GitHub Actions
Error
Written as `cos(n * acos(x))` it routes every value through `acos`, whose derivative is singular at the endpoints. The value and the first derivative survive it, because the clamp zeroes the gradient there, so the forces are fine and the defect is one derivative further out: the second derivative comes back NaN at a pair sitting exactly at the cutoff, which the neighbour list admits, and force training differentiates the force. One such pair poisons the whole gradient. The polynomial has no singularity at the endpoints; the closed form put one there. The recurrence is the same function, smooth to every order, and agrees with the closed form to 2.4e-15 over the interval where the closed form can be evaluated at all. The three bases are now differentiated twice at every distance a pair can legitimately sit at, including the cutoff itself.
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:56 — with
GitHub Actions
Error
GitHub Actions
Error
aacostadiaz
had a problem deploying
to
gpu-internal
September 22, 2026 11:56 — with
GitHub Actions
Error
GitHub Actions
Error
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 1d31017. Configure here.
…ed order Two more from review of the pull request. A model built, moved to a device and then initialised raised a device mismatch. The draw happens on the host deliberately, because a device generator seeded the same way produces different numbers and a recorded seed has to rebuild the same model on any machine, so the drawn tensor is moved to where the parameter is rather than left behind. The contraction hid it: it applies no scale, so its copy crossed devices without complaint while the other two raised. The contraction wrote its canonical weights in storage order, which is output irrep outermost because that is how the forward consumes them. The pinned path order is the enumeration's: body order outermost, output irrep within it. So a mixed-output file held the canonical paths permuted, and a reader splitting it by the counts derived from the basis put the weights on the wrong paths. The round trip through one object could not catch it, because it splits by the counts it wrote itself, so the test that catches it derives them the way another backend would. Both tests fail against the code before this commit.
aacostadiaz
requested review from
ilyes319
and removed request for
ilyes319
September 22, 2026 12:46
aacostadiaz
added
the
stacked
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The kernel backend contract, the five hot ops as torch custom operators, and the plain-torch reference backend every other one is checked against.
What the backend choice stops doing. On the frozen tree it reaches into the model: the interaction blocks branch on whether
conv_fusionhappens to be set at six sites,wrapper_ops.pyreplaces.forwardwith a bound method at three more, and five__new__branches pick a class. Here a descriptor goes in and an op comes out, once, at model build time. Nothing resolves in a forward pass, and the model never learns which backend it got.The op set is closed, because a set that grows to accommodate a backend is not a contract. Five are dispatched:
linear,channelwise_tp_conv,symmetric_contraction,fully_connected_tp,segment_reduce. Two are reference-only and a backend may decline them, which is a normal answer rather than a failure:spherical_harmonicsandradial_basisare cheap closed forms. The span factory that fuses a whole layer returnsNoneand that is the usual case.channelwise_tp_convalways returns node-level values. Whether the reduction over edges is fused into the kernel is the backend's business, which is what deletes theconv_fusionbranch rather than moving it.Compile safety is a property of the ops, not of the caller. Each hot op is a
torch.library.custom_opwith a fake implementation and a registered autograd whose backward is itself differentiable, so force training and double backward are correct. The node count enters as an argument and never asedge_index.max(), so it stays symbolic: a test compiles withfullgraph=Trueacross five batch sizes and asserts one frame compile and no graph break.A backend declares, it never chooses. Capabilities are a coarse filter plus an authoritative
supports(descriptor). A backend without double backward is rejected at build time when forces or stress are being trained, loudly, and stays usable for inference. This is where the Clebsch-Gordan basis stops being decided by what happens to be installed: it is a recorded field, and a backend that cannot serve it fails the build. Nothing readsMACE_USE_CUEQ_CGand nothing infers the basis from whethercuequivarianceimported.One canonical weight layout, which is why the frozen tree needs five conversion command lines and this needs none. The reference holds it directly, so
to_canonicalandload_canonicalare views. The layout carries its normalization folded into the number, where the frozen tree keeps a standard normal and applies a per-path factor inside thee3nnoperation. That makes the factor part of the format, so it is stated with the layout: a linear divides by the square root of the multiplicity feeding its output irrep, the skip connection by the square root of that times the number of element attributes, and the symmetric contraction by nothing.An op that holds weights can also draw a fresh set, and that is not a convenience. Allocating zeros and waiting to be loaded is correct for every model that comes from a checkpoint, and it builds a model that cannot be trained: zero-valued tensors multiply to zero, so does the gradient, and a run started from scratch finishes and reports a model that never moved. No test that loads a checkpoint or overwrites the parameters first can see it. The seed is per op, derived from the run's seed and the op's path through a stable hash, so two ops of the same shape do not receive one draw and a model does not depend on how many random numbers the process had drawn before it was built.
The reference backend uses no
e3nn, asserted by a test, and neither doesmace_core, which imports no framework at all. The symmetric contraction is a Horner cascade over the reduced basis; the spherical harmonics and the radial bases are native closed forms. It is 2,078 lines against the 755 ofwrapper_ops.pyandsymmetric_contraction.pyit replaces, and that is the wrong comparison to draw a conclusion from: those two files dispatch toe3nnandcuequivariancefor the arithmetic, and this one does the arithmetic.What the tests pin:
gradcheckandgradgradcheckat fp64 for the contraction and the convolution and for the whole position-to-force chain; forces that sum to zero; a linear that does not mix one irrep into another and a bias that lands only on scalars; a checkpoint that round trips through a fresh instance; an entry point that resolves; a broken backend that is recorded rather than raised at discovery and quoted when asked for by name; and a freshly built model with no all-zero weighted op.Stacked on #1739, which it needs for the reduced basis, so the commits up to
Convert between the two bases and the two layoutsare that PR's and land with it. Implements BKD-1 (#1560).Note
High Risk
Changes core numerics and checkpoint format (CG basis, weight layout, init scales) and training-critical autograd/compile paths; mistakes would silently alter model size, forces, or portability across backends.
Overview
Introduces a framework-free kernel contract in
mace_coreand a mandatory PyTorch reference backend inmace-torch, so model ops are chosen once at build time from frozen descriptors instead of branching inforwardon optional libraries or fusion flags.mace_core.kernelsadds hashable op descriptors, capability checks (precision, CG basis, double backward), entry-point discovery (mace.kernel_backends.torch), and a pinned canonical checkpoint layout (mul_ir,[Z, A, mul]symmetric-contraction weights, stated init scales).mace_core.clebsch_gordan(with declarednumpy>=1.23) supplies a pure-numpy reduced symmetric tensor-product basis, path labels, and full↔reduced / layout conversions so parameter counts no longer depend on whethercuequivarianceis installed.mace-torchdepends ontorch>=2.4, registersreferencevia entry points, and implements the five dispatched ops astorch.library.custom_opnodes (meta + autograd) plus reference modules for linear, message-passing TP conv, symmetric contraction, skip TP, segment reduce, harmonics, and radial bases. Weight init walks the module tree with stable per-op seeds; canonical save/load is identity for the reference backend.Tests lock basis path counts, oracle agreement (sympy / optional cueq), registry behavior, second derivatives and
torch.compile(fullgraph=True)across batch sizes, and end-to-end force chains through the reference backend.Reviewed by Cursor Bugbot for commit 09d7431. Bugbot is set up for automated code reviews on this repo. Configure here.