Garden's Almanac of Matter Models

TensorNet

January 2023

Checkpoint Params Polaris Sophia Perlmutter Delta Frontier Della
tensornet-matpes-pbe-2025-2 0.8M ○●●○—○
● verified, last 30 days β—‹ installed, not recently verified β€” not installed

No cluster has installed any checkpoint of this model yet.


Running this model
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4with RootstockCalculator(
5 cluster=YOUR_CLUSTER_ID, # eg, "sophia", "perlmutter"
6 checkpoint="tensornet-matpes-pbe-2025-2",
7 device="cuda",
8) as calc:
9 # model is now running in subprocess on compute node
10 atoms.calc = calc
11 atoms.get_potential_energy()

Environments

Rootstock runs each model family inside an isolated Python environment defined by a single file. This environment file includes the specific dependencies needed, plus a setup() function that loads the model and returns an ASE calculator. These files are usually almost identical for a given model family, but because of cluster-specific quirks (eg, an old CUDA driver) the dependencies and setup code can vary a bit.

tensornet_env.py
1# /// script
2# requires-python = ">=3.11"
3# dependencies = [
4# "torch>=2.4.0",
5# "ase>=3.22",
6# "huggingface_hub",
7# "matgl",
8# "nvalchemi-toolkit-ops",
9# "pymatgen",
10# "monty",
11# "ruamel.yaml",
12# "scipy",
13# "torch-geometric",
14# "torch-scatter",
15# "torch-sparse",
16# "torch-cluster",
17# "torch-spline-conv",
18# ]
19#
20# [tool.uv]
21# find-links = ["https://data.pyg.org/whl/torch-2.4.0+cu121.html"]
22#
23# [tool.uv.sources]
24# matgl = { git = "https://github.com/materialsvirtuallab/matgl.git" }
25# ///
26"""TensorNet env β€” hosts MatPES TensorNet checkpoints via MatGL."""
27
28CHECKPOINTS = {
29 "tensornet-matpes-pbe-2025-2": "materialyze/TensorNet-PES-MatPES-PBE-2025.2",
30}
31
32
33def setup(checkpoint: str, device: str = "cuda"):
34 import torch
35
36 torch.set_default_device(device)
37
38 # matgl 1.0.0 imports ExpCellFilter from ase.constraints, but it moved to
39 # ase.filters in ASE 3.23. Patch it in before matgl imports.
40 import ase.constraints
41
42 if not hasattr(ase.constraints, "ExpCellFilter"):
43 from ase.filters import ExpCellFilter
44
45 ase.constraints.ExpCellFilter = ExpCellFilter
46
47 # DGL 2.x graphbolt imports torchdata submodules removed in torchdata>=0.7.
48 # Stub the entire graphbolt subpackage before `import dgl` runs; DGL's
49 # __init__ will use our empty stub and skip the real graphbolt initialisation.
50 # matgl only uses DGL for graph construction β€” graphbolt is never called.
51 import sys, types
52
53 for _name in [
54 "dgl.graphbolt",
55 "dgl.graphbolt.base",
56 "dgl.graphbolt.dataloader",
57 "dgl.graphbolt.feature_fetcher",
58 "dgl.graphbolt.minibatch_transformer",
59 ]:
60 if _name not in sys.modules:
61 sys.modules[_name] = types.ModuleType(_name)
62
63 from huggingface_hub import snapshot_download
64
65 import matgl
66 from matgl.ext.ase import PESCalculator
67
68 # matgl 1.0.0 load_model only checks the GitHub manifest; HF models must
69 # be downloaded explicitly and passed as a local path.
70 local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])
71 pot = matgl.load_model(local_path)
72 return PESCalculator(potential=pot)
73

Built on Polaris: 2026-05-12

Couldn't load the current environments from Rootstock.


References
  1. Simeon, Guillem, de Fabritiis, Gianni, TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials, arXiv, 2023.