deploy.comm

Feature communication and reductions for communication-enabled models.

Outside a communication-enabled variant, gather returns the validated input and reduce returns the rank-local value. During export, both operations lower to FFI calls implemented by chemtrain-deploy and the simulation-engine adapter. LAMMPS provides one such adapter. The local behavior is useful for model initialization and single-domain testing.

The exporter creates the communication object and passes it to energy_fn(). Model implementations use the supplied object at fixed locations in the traced computation. They do not need to construct it themselves.

class ExportCommunication(enabled=False)

Communication interface used while tracing a deployment variant.

gather_widths and reduce_widths record the number of scalar values packed by every gather and reduction during the current trace. A fixed float32 token orders calls through the public FFI. The token value is irrelevant because only the dependency between consecutive calls matters.

Width recording is a Python tracing-time side effect, not part of the compiled computation. The exporter therefore creates a new instance for every trace. Retracing replaces the complete width record instead of appending metadata from an earlier trace.

gather(tree)

Gathers a non-empty, atom-leading floating-point pytree.

All leaves need the same particle dimension and floating-point dtype. Disabled communication still validates the values but does not call the FFI or exchange ghost rows.

reduce(value)

Sums a floating scalar or vector across ranks when enabled.

Non-Communicating Fallback

The module-level function validates and converts a pytree without exchanging data between ranks. It is useful when model setup or reference calculations run outside an exporter-managed communication trace.

gather(tree)

Validates a pytree without communicating between ranks.

Leaves are converted with jax.numpy.asarray(), then returned with the original pytree structure, shapes, and values. Deployment models receive ExportCommunication explicitly. The module-level fallback keeps initialization and single-domain reference calculations executable without a registered native FFI.