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, "polaris", "sophia"
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,<2.5",
5# "ase>=3.22",
6# "huggingface_hub",
7# "matgl",
8# # 0.4+ needs torch>=2.8 at runtime (custom-op registration uses string
9# # annotations infer_schema can't parse on older torch) but only declares
10# # the constraint on its extras, so the resolver won't catch it.
11# # "nvalchemi-toolkit-ops<0.4",
12# "pymatgen",
13# "monty",
14# "ruamel.yaml",
15# "scipy",
16# "torch-geometric",
17# "torch-scatter",
18# "torch-sparse",
19# "torch-cluster",
20# "torch-spline-conv",
21# ]
22#
23# [tool.uv]
24# find-links = ["https://data.pyg.org/whl/torch-2.4.0+cu121.html"]
25#
26# [tool.uv.sources]
27# matgl = { git = "https://github.com/materialsvirtuallab/matgl.git", rev = "35f3ea2d2028febfeb99bb6707f4c6f5223a0851" }
28# ///
29"""TensorNet env — hosts MatPES TensorNet checkpoints via MatGL."""
30
31CHECKPOINTS = {
32 "tensornet-matpes-pbe-2025-2": "materialyze/TensorNet-PES-MatPES-PBE-2025.2",
33}
34
35
36def setup(checkpoint: str, device: str = "cuda"):
37 # matgl 1.0.0 imports ExpCellFilter from ase.constraints, but it moved to
38 # ase.filters in ASE 3.23. Patch it in before matgl imports.
39 import ase.constraints
40
41 if not hasattr(ase.constraints, "ExpCellFilter"):
42 from ase.filters import ExpCellFilter
43
44 ase.constraints.ExpCellFilter = ExpCellFilter
45
46 # DGL 2.x graphbolt imports torchdata submodules removed in torchdata>=0.7.
47 # Stub the entire graphbolt subpackage before `import dgl` runs; DGL's
48 # __init__ will use our empty stub and skip the real graphbolt initialisation.
49 # matgl only uses DGL for graph construction — graphbolt is never called.
50 import sys, types
51
52 for _name in [
53 "dgl.graphbolt",
54 "dgl.graphbolt.base",
55 "dgl.graphbolt.dataloader",
56 "dgl.graphbolt.feature_fetcher",
57 "dgl.graphbolt.minibatch_transformer",
58 ]:
59 if _name not in sys.modules:
60 sys.modules[_name] = types.ModuleType(_name)
61
62 from huggingface_hub import snapshot_download
63
64 import matgl
65 from matgl.ext.ase import PESCalculator
66
67 # matgl 1.0.0 load_model only checks the GitHub manifest; HF models must
68 # be downloaded explicitly and passed as a local path.
69 local_path = snapshot_download(repo_id=CHECKPOINTS[checkpoint])
70 pot = matgl.load_model(local_path).to(device)
71 return PESCalculator(potential=pot)
72
Built on Polaris: 2026-08-03
Couldn't load the current environments from Rootstock.