nvalchemi.distributed.spec.SPEC_LJ_HALO

nvalchemi.distributed.spec.SPEC_LJ_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)

Lennard-Jones pair potential. Halo storage serves cross-rank neighbor pairs from local halo copies after one halo exchange; forces come from direct Warp-kernel writes (no autograd). The wrapper ends its forward with a scatter_add_ aggregating per-atom energies to per-system totals — like MACE’s final scatter_sum — which system_reductions=True routes through per_system_reduce() (slices halo rows off the source and all-reduces).