Metatensor: A Self-Describing Sparse Tensor Data Format for Atomistic Machine Learning

Metatensor is an open-source sparse tensor data format designed specifically for atomistic machine learning. Through its self-describing data structure, native sparse support, and flexible metadata system, it addresses issues faced by traditional tensor libraries when describing atomic systems—such as ambiguous dimensional semantics and difficulties in data exchange—facilitating efficient modeling of atomic systems and cross-team collaboration.

In atomistic machine learning, atomic systems exhibit natural sparsity (only neighboring atoms interact) and complex semantic structures (many atomic properties requiring transformation and combination). Traditional tensor libraries (e.g., NumPy/PyTorch) use anonymous integers for dimensions, lacking physical semantic descriptions and relying on external documentation conventions, which easily leads to errors and collaboration confusion.

1. Self-describing structure: Tensors carry complete metadata, with dimensions as physically meaningful named entities (e.g., atoms, neighbors), allowing structure understanding without external documentation;
2. Native sparse support: Uses sparse storage formats like COO/CSR to reduce memory usage and improve computational efficiency;
3. Flexible metadata system: Supports arbitrary key-value metadata, automatically checks dimensional compatibility during operations to ensure semantic consistency.

• Multi-language and performance: C++ core ensures high performance, with Python bindings provided; supports seamless conversion with frameworks like PyTorch while preserving metadata;
• Simulation software interoperability: Can import data from ASE, LAMMPS, VASP, etc., while preserving metadata such as chemical elements and periodic boundaries;
• ML framework integration: Compatible with graph neural network frameworks like PyTorch Geometric, supporting custom automatic differentiation functions.

• Equivariant neural networks: Annotate transformation properties via metadata to automatically verify operation validity;
• Multi-scale material modeling: Metadata provides a unified cross-scale description, supporting adaptive processing;
• Reproducibility research: Metadata records generation history, facilitating data traceability and collaboration.

1. Expand interoperability with frameworks like JAX and TensorFlow;
2. Introduce distributed computing support (based on Ray/Dask) to handle million-atom systems;
3. Develop domain-specific languages (DSL) to simplify complex model development.

Metatensor drives the evolution of scientific computing data structures toward domain-specific semantic models, enhancing the alignment between code and scientific thinking and lowering the barrier to developing complex models. As its ecosystem matures, it is expected to become a data standard for atomistic machine learning, promoting open collaboration and sustainable research practices in the field.