deepmd_gnn.export

Helpers used by PyTorch-exportable model paths.

Functions

pad_nlist_for_export(→ torch.Tensor)	Append a sentinel neighbor so symbolic export keeps neighbor axes dynamic.
clear_export_guards_once(→ None)	Clear over-specialized guards from the next export of traced.

Module Contents

deepmd_gnn.export.pad_nlist_for_export(nlist: torch.Tensor) → torch.Tensor

Append a sentinel neighbor so symbolic export keeps neighbor axes dynamic.

deepmd_gnn.export.clear_export_guards_once(traced: torch.nn.Module) → None

Clear over-specialized guards from the next export of traced.

make_fx traces may specialize symbolic atom counts too aggressively. DeePMD’s export path calls torch.export.export immediately after tracing, so this one-shot wrapper relaxes those constraints only for that export.