[stefaanhessmann]

Summary

Compiled potentials with their derivatives, plus generative cleanups from sh/v3.

Compile (opt-in everywhere). torch.compile alone cannot train on forces (the force loss needs the gradient of the forces). CompiledPotential records the forward together with the response modules' autograd.grad into one graph, then compiles it:

CompiledPotential(model)        # TypeError: not an AtomisticModel; RuntimeError: torch < 2.14
  __call__(inputs)
    key = (train/eval, postprocessing, input keys/dtypes/devices, sizes 0 or 1)
    once per key: make_fx(forward incl. forces/stress) -> torch.compile(dynamic=True)
    train: graph(params, buffers, inputs)   # force loss needs one backward
    eval:  graph under no_grad

where	option (default off)
AtomisticTask / spktrain	task.compile=True (distillation runs uncompiled, with a warning)
spkpredict	compile=True
spkmd SchNetPackCalculator	calculator.compile_model=True
ASE SpkCalculator	compile_model=True
spk.dynamics.Calculator	compile=True

Coulomb, Ewald and ZBL count molecules from n_atoms instead of int(idx_m[-1]) + 1 (a GPU sync and a data-dependent size).

Generative.

• e8855064: the sampler's calculator only runs the model (JointCalculator renamed FieldCalculator); updates convert their head and add guidance. GenerativeCalculator, the Score/Noise/Velocity/X0/PseudoForce calculators and ParametrizedModel are removed.
• 05903841: ReverseODE removed, differential_equations renamed sde.
• 688e6cb2: MatchingLoss removed; training goes through Diffuse.
• Merge of v3 (040aca4f): the new docs/api/generative.rst drops the removed names.

Evidence

• tests/atomistic/test_compile.py (float64, default assert_close), all marked slow. These checks passed on 2026-10-26 (laptop GPU, torch 2.14.1); they have not been rerun since the last commit (2f2a1033):
  • SchNet, PaiNN, SO3net: compiled energy, forces and force-loss parameter gradients equal eager, in train and eval mode.
  • Stress on periodic silicon cells.
  • One-molecule batches get their own graph.
  • AtomisticTask.training_step gradients equal eager; the task pickles.
  • The three calculators give eager forces.
• Default suite on 2f2a1033: 788 passed, 0 failed, 13 skipped (12 need mace-torch, 1 by design).
• Compiled vs eager runs:
  • spktrain experiment=md17, synthetic 80-molecule db, 3 epochs: Before (eager) val_loss 190646 / 161211 / 152163, test_forces_mae 332.089; After (task.compile=True) identical to 6 digits.
  • spkpredict compiled vs eager: max relative difference 2e-8 for energy and 3e-7 for forces (float32).
• Speed (earlier probes, H100 PCIe, SchNet 6 × 128, 21 to 8192 atoms): compiled vs eager 1.8–3.1× for energy+forces, 1.7–2.2× for a training step; about 10–25 s to compile on the first step. A timed end-to-end QM9 training run is still to do.

Merge Danger

Door: two-way

Compile is off by default, and the code paths are unchanged unless an option is set.

Blast Radius: generative API

Code that uses the removed generative classes breaks: MatchingLoss, ReverseODE, ParametrizedModel, GenerativeCalculator, JointCalculator and the per-field score calculators. The new compile options need torch ≥ 2.14 when switched on; the package floor stays at torch 2.5.

🤖 Generated with Claude Code