[BUG] MemmapDataset crashes on target variants (energy/<variant>) + non_conservative_stress/<variant> dangerously skips process_non_conservative_stress()

[zyxwwxyz]

Summary

Training on a MemmapDataset with a variant target (e.g. energy/pbe) fails at the first batch with:

invalid parameter: 'energy/pbe' is not a valid label name

The identical target definition trains fine when the data is read from an ASE-readable .xyz file. The issue should be isolated to MemmapDataset.__getitem__, which builds the metatensor property Labels from the raw target name and does not strip the /<variant> suffix.

Is this a wanted behaviour or a bug worth correcting?

Expected behavior

Training runs correctly with /<variant> suffix

Actual behavior

mtt stops with error

invalid parameter: 'energy/pbe' is not a valid label name

Version

2026.2.1

Steps to reproduce

from metatensor.torch import Labels

# What `MemmapDataset.__getitem__` currently does:

Labels.range("energy/pbeu".replace("mtt::", ""), 1)
-> RuntimeError: invalid parameter: 'energy/pbeu' is not a valid label name

# What `target_info.py` / the ASE reader do (works):

tk = "energy/pbeu"
Labels.range((tk.split("/")[0] if "/" in tk else tk).replace("mtt::", ""), 1)  # OK -> "energy"

Further information, files, and links

Suggested fix:

properties=Labels.range(
-            target_key.replace("mtt::", ""), target_array.shape[-1]
+            (target_key.split("/")[0] if "/" in target_key else target_key).replace(
+                "mtt::", ""
+            ),
+            target_array.shape[-1],
         ),

[Luthaf]

Is this a wanted behaviour or a bug worth correcting?

This is a bug for sure, the expected data format is clear: https://docs.metatensor.org/metatomic/latest/quantities/energy.html. The property name should be energy regardless of the variant.

[zyxwwxyz]

Ok. Shall I send a PR with the suggested fix then?

[Luthaf]

I'm not fully sure the suggested fix will work in all possible cases, the name of the property dimension and the name of the target don't have to match. Maybe adding a specific function property_name_for_target or whatever that can do the right thing for all standard targets, and then do whatever is needed for custom mtt outputs. The suggested fix would work fine for the energy though.

@pfebrer, @PicoCentauri, @frostedoyster any preference on the way to fix this?

[pfebrer]

I have no idea how variants or memmaps work 😅

[PicoCentauri]

I have no strong feeling on how to fix this. Where should the property_name_for_target function live @Luthaf?

[Luthaf]

Wherever convenient in metatrain was my thinking

[zyxwwxyz]

Coming back to this issue: as a workaround I trained my model with the bare energy label, loaded the final .ckpt file, and edited the head name to energy/variant by hand. It works fine for all predictions (energies and conservative forces and stresses).
However, I need to do the same also for a non_conservative_stress head, so I changed its name to non_conservative_stress/variant by hand in a similar fashion.
However, this silently breaks the TODO at

metatrain/src/metatrain/pet/modules/backend.py

Lines 464 to 471 in 12da07e

	if output_name == "non_conservative_stress": # TODO: variants
	num_properties = block_sums[0].shape[1] // 9
	block_sums[0] = process_non_conservative_stress(
	block_sums[0],
	cells,
	system_indices,
	num_properties,
	)

and returns a non-reshaped, non-normalized stress tensor. I think it is useful to remind this when solving this bug and maybe find a common solution.

[jwa7]

Is this closed by #1216 @zyxwwxyz ?

[zyxwwxyz]

Is this closed by #1216 @zyxwwxyz ?

Yes it looks like #1216 closes this. However, I think my last comment about non_conservative_stress/variant skipping process_non_conservative_stress is still open. Should I open a new, separate issue for that?

[jwa7]

Ah no in that case this issue can stay open, thanks :)