Python operations reference

The Python API for metatensor also provides functions which operate on metatensor.TensorMap, metatensor.TensorBlock, and metatensor.Labels and can be used to build your own Machine Learning models.

The list of all operations currently implemented is available below. If you need any other operation, please open an issue!

Using operations with PyTorch

PyTorch is a very popular framework for machine learning, providing multiple tools to make writing and training models easier. There are two ways to use the operations with PyTorch:

• Using the pure Python version of metatensor, one can store values in a metatensor.TensorBlock using torch.Tensor. In this case, all operations will be compatible with torch autograd (automatic gradient tracking and differentiation). This allows to train models from Python, but not to export the models to run without the Python interpreter. When running a model with the pure Python version of metatensor, you should use the operations from metatensor.<operation_name>.

• When using the TorchScript version of metatensor, one can also compile the Python code to TorchScript and then run the model without a Python interpreter. This is particularly useful to export and then use an already trained model, for example to run molecular simulations. If you want to do this, you should use classes and operations from metatensor.torch, i.e. metatensor.torch.TensorMap and using the operation from metatensor.torch.<operation_name>. All the operation above are available in the metatensor.torch module, but some might not yet be fully TorchScript compatible.

List of all operations

• Checks
• Creation operations
  • empty_like()
  • ones_like()
  • zeros_like()
  • random_like()
  • block_from_array()
• Linear algebra
  • dot()
  • lstsq()
  • solve()
• Logic function
  • allclose()
  • equal()
  • equal_metadata()
• Manipulation operations
  • drop_blocks()
  • join()
  • manipulate dimension
  • one_hot()
  • remove_gradients()
  • samples reduction
  • slice()
  • split()
  • to()
• Mathematical functions
  • abs()
  • add()
  • divide()
  • multiply()
  • pow()
  • subtract()
• Set operations
  • unique_metadata()
  • sort()