xnn.common.deploy.torchscript.split_wrappers#

xnn.common.deploy.torchscript.split_wrappers(model, total_charge=0.0)[source]#

Peel the LES / D4 wrappers off a model.

Parameters:

model (torch.nn.Module) – A core model, optionally wrapped in LatentEwald and/or D4Dispersion (any nesting order).
total_charge (float, optional) – Net charge handed to the D4 head, by default 0.

Returns:

core (torch.nn.Module) – The innermost model (a _ZeroCore for a standalone D4).
long_range (torch.nn.Module) – _LatentEwaldHead or _NoLongRange.
dispersion (torch.nn.Module) – _DispersionHead or _NoDispersion.
core_cutoff (float) – The core model’s neighbor-list radius (edges beyond it are filtered out before the core runs).

Raises:

TypeError – If LES sits directly on a standalone D4 model (its features are the D4 per-atom quantities, which the tensor ABI does not reproduce).