nvalchemi.distributed.MethodAdapter

class nvalchemi.distributed.MethodAdapter(target=None, method_or_class=None, replacement_or_method=None, replacement=None, *, module_path=None, class_name=None, method_name=None, mode='wrap', install_site='')

Wrap a class method: intercept the call, transform an argument, then invoke the original — as opposed to PythonAdapter / JitAdapter, which replace a module-level function outright.

Some third-party models need a method-internal arg wrap that can’t be expressed by replacing a module-level helper. Canonical case: AIMNet2’s aimnet.modules.aev.ConvSV.forward — its conv_q path (d2features=False) indexes a charges-derived arg that has lost ShardTensor metadata through the MLP, so it must be re-wrapped as a sharded ShardTensor before the stock a.index_select(0, nbmat.flatten()), or it reads global nbmat indices off a rank-local tensor.

replacement is a wrapping function (original, *args, **kwargs): install() captures the original (unbound) method and binds it as the first argument, so the replacement can transform args and call through. Reading per-step routing off the method’s own arguments (e.g. ConvSV’s data dict, which AIMNet2Wrapper.adapt_input() populates) keeps it free of ambient context — the same discipline the PythonAdapter replacements follow. Declared on a spec’s third_party_helpers so the framework’s AdapterRegistry installs + restores it; no wrapper distributed_setup hook needed.

Normally name the real imported class, not a string path:

MethodAdapter(ConvSV, "forward", _rewrap_conv_q)   # real class + method

The string form MethodAdapter("module", "Class", "method", ...) is also accepted (disambiguated by the first argument’s type), so serialized dicts round-trip unchanged.

Parameters:

• target (type | str | None)

• method_or_class (str | None)

• replacement_or_method (Any)

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

• module_path (str)

• class_name (str)

• method_name (str)

• mode (Literal['wrap', '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 MethodAdapter from to_dict() output.

Parameters:

d (dict[str, Any])

Return type:

MethodAdapter

install()

Wrap class.method so calls route through replacement.

replacement=None is the declaration-only form (registry no-ops); present for parity with the other adapters.

mode="marshal": wrap the whole method (e.g. e3nn SphericalHarmonics.forward) with make_marshaller() — the smallest region that resolves BOTH (A) a scripted call inside the method AND (B) a subsequent in-place mutation of a ShardTensor (e.g. sh.mul_(cat)): the marshaller unwraps the ShardTensor input to its local ONCE, so the scripted op and the in-place op both run on a plain local tensor (eager, torch.compiler.disable graph-break), then the output is re-wrapped. Subsumes a separate JitAdapter on the inner scripted function. (B is a PyTorch AOT limitation — in-place mutation of a subclass that is a graph input across a graph break — reproduced on stock TwoTensor; keeping the region eager sidesteps it.)

Return type:

dict[str, Any]

restore(memento)

Reverse install(): put the original method back.

Parameters:

memento (dict[str, Any])

Return type:

None

to_dict()

Serialize to a JSON-roundtrippable dict.

Return type:

dict[str, Any]