nvalchemi.distributed.PythonAdapter

class nvalchemi.distributed.PythonAdapter(module_path, attr_name, replacement=None, install_site='')

Replace a plain-Python module-level helper with a distributed-aware version.

Unlike JitAdapter, the helper isn’t @torch.jit.script — it’s a normal Python function whose single-process tensor-layout assumptions break under partition. Canonical case: aimnet.nbops.mol_sum reading mol_idx[-1] + 1 for its output size.

replacement may be None if it must be built at install time by the wrapper (e.g. it closes over per-run partition metadata that’s only available in DistributedModel.__enter__). Pass a factory via the wrapper’s distributed_setup hook or override install() in a subclass.

Parameters:

• module_path (str)

• attr_name (str)

• replacement (Callable[[...], Any] | None)

• install_site (str)

describe(state='pending', error=None)

Return an AdapterStatus snapshot of this adapter.

Parameters:

• state (Literal['pending', 'installed', 'restored', 'failed'])

• error (str | None)

Return type:

AdapterStatus

classmethod from_dict(d)

Reconstruct a PythonAdapter from to_dict() output.

Parameters:

d (dict[str, Any])

Return type:

PythonAdapter

install()

Swap in the plain-Python replacement at module.attr.

replacement=None = declaration-only: the wrapper’s distributed_setup is responsible for swapping the attr. See note on JitAdapter.install() for the rationale.

Return type:

dict[str, Any]

restore(memento)

Reverse install(): put the original attribute back.

Parameters:

memento (dict[str, Any])

Return type:

None

to_dict()

Serialize to a JSON-roundtrippable dict.

Return type:

dict[str, Any]