nvalchemi.distributed.spec.SPEC_PME_HALO

nvalchemi.distributed.spec.SPEC_PME_HALO = MLIPSpec(distribution=DistributionSpec(policy=HaloStoragePolicy(scatter_mode='halo_correction', gather_mode='halo_read'), custom_ops=(), third_party_helpers=(), adapters=(), shard_fields=('positions', 'charges', 'atomic_numbers')), owned_only_outputs=frozenset(), all_reduce_outputs=frozenset(), output_kinds={'energy': <OutputKind.PER_GRAPH: 'per_graph'>, 'forces': <OutputKind.PER_NODE: 'per_node'>, 'stress': <OutputKind.PER_GRAPH: 'per_graph'>, 'atomic_energies': <OutputKind.PER_NODE: 'per_node'>}, system_reductions=True, node_energy_key=None, node_virial_key=None, gp_replicate_geometry=False, outputs=None, compile=None)

halo storage. Real-space pair interactions on halo-padded inputs follow the standard halo path; charge spreading gets an owned-slice + all-reduce handler. Post-spread stages (FFT, Green’s function, IFFT, spline_gather, corrections) are replicated across ranks — they operate on the all-reduced global mesh, so no dispatch is needed. Caveat: single-system halo (batch_idx=None) hits a plain charges.sum() in pme_energy_corrections that double-counts halo rows; batched halo works correctly via scatter_add_ dispatch through per_system_reduce().

Type:

PME (Particle Mesh Ewald)