fix: handle missing gradients in remove_additive for mixed datasets

[ChaitanyaParate]

Fixes #867

Root cause

When training on mixed datasets (some structures with stress, some without), the additive model does not compute strain gradients for samples from the no-stress subset. _add_block_block requires structural parity between both blocks, so it crashed with RuntimeError: invalid parameter: can not find gradients with respect to 'strain' in this block.

Fix

When a gradient exists in the target block but not in the additive block, a zero-valued placeholder is inserted (same shape and metadata as the target gradient). The additive model's contribution to that gradient is zero, so subtracting zero leaves the target unchanged. A warnings.warn is emitted so the user is aware a placeholder was inserted.

Additional fix

The .to(device=device) call was missing in the else branch for gradient values. This is corrected here.

Self-introduced regression corrected

An earlier iteration of this PR accidentally switched the else branch to use target_gradient.components/properties. Git history confirms the original code always used additive_gradient.components/properties with only samples coming from the target. This is restored.

Testing

A unit test test_remove_additive_missing_gradient is added to tests/utils/test_additive.py. It uses a stub model that guarantees no strain gradient in the additive output, verifies the warning fires, and checks the target gradient values are unchanged after subtraction.

Note: I did not test on the MAD dataset directly. The test uses a synthetic TensorMap that reproduces the structural mismatch. Maintainer verification on real data would be appreciated.

---

📚 Documentation preview 📚: https://metatrain--1140.org.readthedocs.build/en/1140/

[pfebrer]

Hi, since I don't use gradients I don't fully understand what this PR solves, could you provide a zip file with an example that would fail before this PR but works once this is implemented? Thanks!

[ChaitanyaParate]

Hi, since I don't use gradients I don't fully understand what this PR solves, could you provide a zip file with an example that would fail before this PR but works once this is implemented? Thanks!

Hi! The zip contains a self-contained script (reproduce_bug.py) that demonstrates the crash and the fix with no dataset or model checkpoint needed.

Setup:

pip install metatrain metatensor-operations
python reproduce_bug.py
Part 1 directly calls the same _add_block_block metatensor operation shown in the original traceback, with the exact gradient mismatch — you'll see:

RuntimeError: invalid parameter: can not find gradients with respect to 'strain' in this block
Part 2 applies the fix (zero placeholder) and re-runs the same call — all assertions pass.

example_mixed_stress_bug_repro.zip

[pfebrer]

Thanks, actually I wanted an example with a dataset and an options file to understand when is this a problem in real life. It is hard for me to understand how the script relates to a real-life case. Is this a problem that you are facing when trying to use metatrain?

[ChaitanyaParate]

Hi! I attempted to reproduce it end-to-end with a real dataset and options.yaml, and I want to be transparent about what I found.

The real-life scenario is: training with two datasets simultaneously, one with stress and one without:

yaml
training_set:

• systems: {read_from: carbon_with_stress.xyz}
  targets:
  energy:
  virial: {read_from: carbon_with_stress.xyz, key: virial}
  forces: ...
• systems: {read_from: carbon_no_stress.xyz}
  targets:
  energy:
  forces: ... # no virial/stress here
  Why I can't reproduce it end-to-end today: Between version 2025.11 (where the issue was reported) and the current main, evaluate_model was independently updated so that it now always injects autograd strain gradients for energy targets regardless of what the model returns. This means remove_additive no longer sees the missing gradient through the normal pipeline.

However, the defensive fix in remove_additive is still valid and valuable — any future additive model that genuinely doesn't produce a gradient (e.g. a non-energy additive, or a model returning a pre-computed output without grad_fn) would hit this path.

The reproduce_bug.py in the zip demonstrates the exact failing metatensor call (_add_block_block) from the original traceback with the exact error message — that part is real and version-independent.

example_mixed_stress_bug_repro.zip