nvalchemi.distributed.CompilePolicy

class nvalchemi.distributed.CompilePolicy(static_shapes=True, graph_padder=None, force_strategy=ForceStrategy.MODEL_INTERNAL, stress_via_strain=False)

How a model wants torch.compile driven under domain decomposition.

static_shapes requests fixed-shape (capped) compilation so a compiled MD trajectory stays compiled across steps — the framework pads the graph to stable per-rank capacities. graph_padder is the GraphPadder used for that padding; when None the framework uses the built-in COO edge_index padder (COOPadder), so a standard MPNN declares nothing. force_strategy declares how forces are produced (see ForceStrategy); the derived forces_via_autograd / consolidate_node_energy / energy_output properties follow from it.

The policy is only the contract. Whether and how to compile is owned by DistributedModel (constructed with compile=True / compile_kwargs=...); the policy carries no compile switch of its own.

Parameters:

• static_shapes (bool)

• graph_padder (GraphPadder | None)

• force_strategy (ForceStrategy)

• stress_via_strain (bool)

property consolidate_node_energy: bool

True when the model returns un-reduced per-node energy and the framework does the owned-only per-graph sum + all-reduce (FRAMEWORK_FROM_NODE_ENERGY).

property energy_output: str

The active_outputs key driven for the energy-only forward: "atomic_energies" for the per-node strategy, else "energy".

property forces_via_autograd: bool

True when the framework owns the energy-only forward + force autograd (any FRAMEWORK_FROM_* strategy).