Metatomic interface seems not torch scriptable

[PicoCentauri]

When trying to export a simple test function

from metatomic.torch import AtomisticModel, System
import vesin.metatomic

@torch.jit.script
def foo(systems: list[System], model: AtomisticModel) -> None:
    vesin.metatomic.compute_requested_neighbors(
        systems=systems, system_length_unit="nm", model=model
        )

I get an error for an unsupported annotation

AssertionError: Unsupported annotation typing.Union[metatomic.torch.model.AtomisticModel, metatomic.torch.model.ModelInterface] could not be resolved because None could not be resolved. At
  File "/Users/philiploche/projects/water-benchmark/.venv/lib/python3.13/site-packages/vesin/metatomic/_model.py", line 17
    systems: Union[List[System], System],
    system_length_unit: str,
    model: Union[AtomisticModel, ModelInterface],
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    model_length_unit: Optional[str] = None,
    check_consistency: bool = False,

[Luthaf]

Yes, for now only vesin.metatomic.NeighborList is compatible with TorchScript, when initialized with torchscript=True.

Where do you need to use compute_requested_neighbors inside TorchScript?

[PicoCentauri]

I am writing a wrapper for the GROMACS native MLIP interface to benchmark it against our own interface. Would be just convenient to use the compute_requested_neighbors:

class NPotModelWrapper(torch.nn.Module):
    def __init__(self, model: AtomisticModel):
        super().__init__()
        self.model = model

    def forward(
        self,
        positions,
        atomic_numbers,
        box: Optional[torch.Tensor] = None,
        pbc: Optional[torch.Tensor] = None,
    ):
        if box is None:
            box = torch.zeros((3, 3), dtype=positions.dtype, device=positions.device)

        if pbc is None:
            pbc = torch.ones((3,), dtype=torch.bool, device=positions.device)

        systems = [System(types=atomic_numbers, positions=positions, cell=box, pbc=pbc)]

        vesin.metatomic.compute_requested_neighbors(
            systems=systems,
            system_length_unit="nm",
            model=self.model,
        )

        results = self.model.forward(
            systems, options=self._options, check_consistency=False
        )

        return results["energy"].block(0).values[0][0]