[BUG] Positions and cell gradients instability in autograd tests

[abmazitov]

Summary

As discussed and discovered in #959 , it looks like the current settings in the autograd tests are somewhat "magical". They are using a system defined like this:

system = System(
    types=torch.tensor([6, 6], device=device),
    positions=torch.tensor(
        [[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]],
        dtype=torch.float64,
        device=device,
        requires_grad=True,
     ),
     cell=torch.eye(
        3, dtype=torch.float64, requires_grad=True, device=device
     ),
     pbc=torch.tensor([True, True, True], device=device),
)

When using PET with a cosine cutoff and a width of 0.5, @ceriottm and I noticed that the test is failing. Our original idea was that the cutoff function has a discontinuity at the location of atoms and / or box size, and therefore decided to change the system parameters a bit.

A tiny change in the cell definition from torch.eye(3) to 1.1 * torch.eye(3) in the test_autograd_cell however have resulted in this test failing at least for SOAP-BPNN.

Expected behavior

Gradcheck and gradgradcheck passing irrespective of the reasonably meaningful input parameters

Actual behavior

In the current state, only a certain magical configuration mentioned above was making the tests work.

Version

2025.13.dev68+gbb5dacf3e.d20251210

Steps to reproduce

Change the cell to

cell = 1.1 * torch.eye(3, dtype=torch.float64, requires_grad=True, device=device)

in the line 136 in the metatrain/utils/testing/autograd.py and run the soap-bonn tests

Further information, files, and links

No response

[ceriottm]

My impression is that the numerical precision of gradgrad is not great for many architectures, which is what leads to these failures. For instance soap_bpnn can be made to work with the standard tolerances by reducing nmax,lmax and/or the number of layers. I would be inclined to keep the higher tolerance as we want to test realistic use cases

[pfebrer]

Yes, although I would at some point take the time to find the minimum tolerance for each architecture and set it for the respective tests, in this way we can easily see if a change in the architecture makes gradients worse (and perhaps simply increase the tolerance, but at least we notice the change).