Create a System class in Python by EricBoittier · Pull Request #325 · metatensor/metatomic

EricBoittier · 2026-09-15T07:32:05Z

Contributor (creator of pull-request) checklist

Tests updated (for new features and bugfixes)?
Documentation updated (for new features)?
Issue referenced (for PRs that solve an issue)?

Reviewer checklist

CHANGELOG updated with public API or any other important changes?

EricBoittier

👀

Luthaf

Did you check for compatibility with torch.compile/jax.jit?

EricBoittier · 2026-09-17T09:39:29Z

Did you check for compatibility with torch.compile/jax.jit?

I generated a notebook to answer these questions (attached as PDF, sorry for LLM-braining
system-compile-jit.pdf
)

It shows:

numpy / torch / jax construction and getters work
extracted arrays work with torch.compile (fullgraph=True) and jax.jit
constructing System inside torch.compile(..., fullgraph=True) fails on ctypes POINTER.new and ctypes graph-breaks to eager if fullgraph!=True
jax.jit rejects System; constructing inside jit fails (dlpack) since metatomic.System is a host-side ctypes wrapper around mta_system_t, not a traced model type.

Luthaf · 2026-09-17T12:41:12Z

I'd say it's fine for now if we can not construct a system inside a torch.compile/jax.jit bit of code, but it would be nice to be able to compile a function like squared_sum_positions(system): return torch.sum(system.positions)**2. Looks like this works for torch, how hard would it be to make it work for jax? We don't have to do it in this PR, just make sure we are not painting ourself in a corner!

Luthaf

I do not understand the new requirement of releasing all arrays views before adding pairs/custom data. Can you explain?

EricBoittier · 2026-09-23T09:18:49Z

I do not understand the new requirement of releasing all arrays views before adding pairs/custom data. Can you explain?

The two APIs do not share buffers.
add_pairs only does system_mut.pairs.insert(...) (system.rs 371).
They're coupled because mta_system_get_data clones the Arc into manager_ctx (c_api/system.rs 191–209). add_pairs / add_custom_data go through Arc::get_mut (281–285 / 408–412), so any live getter view makes strong_count > 1, so mutation fails

We can keep this, or allow adds with live views (interior mutability on pairs / custom_data only). Which do you want?

On the topic of mutability, jax arrays strictly immutable. In torch, arrays may still look writable depending on version; the C comment on get_data is still “do not modify” (c_api/system.rs 168–170).
I would go for immutable, even if pytorch doesn't strictly require it.

Luthaf · 2026-09-23T12:42:43Z

Ah, right there is the Arc::get_mut thing. Ok, makes sense then!

On the topic of mutability, jax arrays strictly immutable. [...]

This should already be handled by DLPack, which has a flag for immutable data. So unless we gain something by specifying it ourself, I'd leave read-only considerations to DLPack and not mention it ourself.

Luthaf · 2026-09-24T12:02:30Z

+    if module.startswith("torch") and name == "Tensor":
+        return "torch"
+    if module.startswith("jax") or module.startswith("jaxlib"):
+        return "jax"


should maybe use isinstance instead since we could have subclasses of these. See https://github.com/metatensor/metatensor/blob/9dfa936675f8ee92087ca63d122ac04f76333447/python/metatensor_operations/metatensor_operations/_dispatch.py#L44

Still relevant

Luthaf · 2026-09-24T12:03:03Z

+    raise ValueError(f"Unknown arrays backend: {backend}")
+
+
+def _pair_options_json(options: PairListOptions) -> bytes:


not worth a helper IMO

Luthaf · 2026-09-24T12:05:12Z

+        The returned array uses the configured :py:attr:`arrays_backend` and is
+        a read-only view.


Suggested change

The returned array uses the configured :py:attr:`arrays_backend` and is

a read-only view.

The returned array uses the configured :py:attr:`arrays_backend`.

and same below

Luthaf · 2026-10-01T13:53:47Z

I rebased and squashed all commits into a single one, but there are still some unresolved review comments here.

to pull the rebased changes:

git checkout python-wrap-system
git fetch
git reset --hard origin/python-wrap-system # (assuming the remote it called origin)

Co-authored-by: Cursor <cursoragent@cursor.com>

Added comments to clarify the use of `ManuallyDrop` in `mta_system_add_pairs` and `mta_system_add_custom_data` functions, ensuring that the caller's system remains valid even when these functions fail. Introduced a new test case to verify that the system is not dropped on failure during pair and custom data additions.

EricBoittier changed the base branch from main to metatomic-core September 15, 2026 07:32

Luthaf changed the title ~~Add an empty metatomic-core python package, re-exporting metatomic-torch~~ Create a System class in Python Sep 15, 2026

EricBoittier force-pushed the python-wrap-system branch from 53e7927 to bca7450 Compare September 15, 2026 11:13

EricBoittier marked this pull request as ready for review September 15, 2026 11:15

Luthaf reviewed Sep 15, 2026

View reviewed changes

EricBoittier commented Sep 15, 2026

View reviewed changes

EricBoittier requested a review from Luthaf September 15, 2026 13:23

Luthaf reviewed Sep 16, 2026

View reviewed changes

EricBoittier force-pushed the python-wrap-system branch from 67d8cd2 to 71002d0 Compare September 17, 2026 09:54

EricBoittier requested a review from Luthaf September 18, 2026 09:48

Luthaf reviewed Sep 23, 2026

View reviewed changes

Luthaf reviewed Sep 24, 2026

View reviewed changes

Luthaf force-pushed the metatomic-core branch 2 times, most recently from 9b44483 to 15fd62d Compare September 29, 2026 15:12

EricBoittier force-pushed the python-wrap-system branch from 4773782 to 6c05848 Compare September 30, 2026 07:33

Luthaf force-pushed the metatomic-core branch 2 times, most recently from 5ec31c2 to 09b9cf2 Compare October 1, 2026 10:54

EricBoittier requested a review from Luthaf October 1, 2026 13:24

Luthaf force-pushed the python-wrap-system branch 2 times, most recently from 38c8afa to 68614de Compare October 1, 2026 13:45

Add a Python System class wrapping mta_system_t

33fd484

Luthaf force-pushed the python-wrap-system branch from 68614de to 33fd484 Compare October 1, 2026 13:51

EricBoittier and others added 2 commits October 1, 2026 16:36

Address review comments on the Python System class.

aff9d7d

Co-authored-by: Cursor <cursoragent@cursor.com>

		raise ValueError(f"Unknown arrays backend: {backend}")


		def _pair_options_json(options: PairListOptions) -> bytes:

		The returned array uses the configured :py:attr:`arrays_backend` and is
		a read-only view.

Conversation

EricBoittier commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Contributor (creator of pull-request) checklist

Reviewer checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

EricBoittier left a comment

Choose a reason for hiding this comment

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

EricBoittier commented Sep 17, 2026

Uh oh!

Luthaf commented Sep 17, 2026

Uh oh!

Luthaf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

EricBoittier commented Sep 23, 2026

Uh oh!

Luthaf commented Sep 23, 2026

Uh oh!

Uh oh!

Luthaf Sep 24, 2026

Choose a reason for hiding this comment

Uh oh!

Luthaf Oct 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Luthaf Sep 24, 2026

Choose a reason for hiding this comment

Uh oh!

Luthaf Sep 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Luthaf commented Oct 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

EricBoittier commented Sep 15, 2026 •

edited

Loading