API reference¶

The classes and functions in the TorchScript API are kept as close as possible to the classes and functions of the pure Python API, with the explicit goal that changing from

import metatensor
from metatensor import TensorMap, TensorBlock, Labels


import metatensor.torch as metatensor
from metatensor.torch import TensorMap, TensorBlock, Labels

should be 80% of the work required to make a model developed in Python with metatensor.operations compatible with TorchScript. In particular, all the operations are also available in the metatensor.torch module under the same name. All the functions have the same behavior, but the versions in metatensor.torch are annotated with the types from metatensor.torch, and compatible with TorchScript compilation. For example metatensor.add() is available as metatensor.torch.add().

The learn module is also re-exported inside metatensor.torch.learn, with the same functionalities as metatensor.learn.

The documentation for the usual core classes of metatensor can be found in the following pages: