nvalchemi.distributed.JitAdapter

class nvalchemi.distributed.JitAdapter(module_path, attr_name, replacement=None, mode='eager', install_site='')

Replace a @torch.jit.script-decorated module-level helper so a ShardTensor can cross it safely on the distributed path.

Two modes:

• mode="marshal": wrap the original scripted op with make_marshaller() at install time — unwrap ShardTensor→local, run the scripted op, rewrap. Keeps the op scripted/fused; no hand-written copy.

• mode="eager": swap in replacement (a hand-written plain-Python equivalent so ShardTensor __torch_function__ fires inside it). The author keeps that copy in sync with upstream.

Parameters:

• module_path (str)

• attr_name (str)

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

• mode (Literal['eager', 'marshal'])

• 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 JitAdapter from to_dict() output.

Parameters:

d (dict[str, Any])

Return type:

JitAdapter

install()

Swap in the replacement at module.attr.

mode="marshal": build a make_marshaller() around the current attribute (the original scripted op) — no hand-written copy needed.

mode="eager": swap in self.replacement. replacement=None is the declaration-only form: the entry is in the spec for diagnostics, but the wrapper’s distributed_setup hook swaps the attribute (it closes over per-run partition metadata). The registry no-ops here.

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]