nvalchemi.distributed.DistributionSpec

class nvalchemi.distributed.DistributionSpec(policy=None, custom_ops=<factory>, third_party_helpers=<factory>, adapters=<factory>, shard_fields=('positions', 'charges', 'atomic_numbers'))

Framework-generic distributed spec.

Carries the field’s StoragePolicy (how local storage relates to its placement + the overlay-aware op behavior) and the declarative tuples of third-party touchpoints. No chemistry vocabulary appears here — output names like "stress" / "forces" live in MLIPSpec one layer up.

Parameters:

• policy (nvalchemi.distributed._core.storage_policy.StoragePolicy | None) – A StoragePolicy (HaloStoragePolicy / PlainShard), or None for the local (no cross-rank) case. The dispatch attaches it to each tensor as _storage_policy and routes ops through it.

• adapters (tuple[Any, ...]) – The single declarative field for registering adapters — one tuple mixing OpAdapter (opaque custom/triton kernels) with JitAdapter / PythonAdapter / MethodAdapter (third-party callable replacements). At construction it is lowered: each OpAdapter is appended to custom_ops, everything else to third_party_helpers, and adapters is cleared. The two split tuples remain the canonical storage all framework consumers read, so serialization and dispatch are unchanged — adapters is purely a unifying constructor convenience.

• custom_ops (tuple[nvalchemi.distributed._core.adapter.OpAdapter, ...]) – Tuple of OpAdapter declaring custom-op wrap config (kernels registered via @torch.library.custom_op / @torch.library.triton_op). May be passed directly or supplied via adapters.

• third_party_helpers (tuple[Any, ...]) – Tuple of JitAdapter / PythonAdapter / MethodAdapter — third-party callables that need a distributed-aware replacement. May be passed directly or via adapters.

• shard_fields (tuple[str, ...])

classmethod from_dict(d)

Inverse of to_dict().

Parameters:

d (dict[str, Any])

Return type:

DistributionSpec

to_dict()

JSON-friendly representation. Used by MLIPSpec.to_dict() (v2 schema, nested under "core").

Return type:

dict[str, Any]