deploy.exporter

Exporting potential models to serialized StableHLO bundles.

class ParticleField(name: str, dtype: Any = <class 'jax.numpy.int32'>)

Descriptor for an additional scalar per-particle model input.

species is implicit and must not be registered here. Additional field semantics are defined by the exported model and its accompanying documentation.

dtype: Any

Alias for field number 1

name: str

Alias for field number 0

class GlobalField(name: str, dtype: Any = <class 'jax.numpy.float32'>)

Descriptor for a scalar model input shared by the full configuration.

dtype: Any

Alias for field number 1

name: str

Alias for field number 0

class OutputScope(value)

Declares the leading shape and engine reduction of an output.

PARTICLE outputs have a leading particle axis. LOCAL outputs are additive rank-local configuration values that the engine adapter reduces once. GLOBAL outputs are complete configuration values and must not be reduced again.

class OutputField(name: str, scope: OutputScope = OutputScope.PARTICLE, extensive: bool = False)

Descriptor for an auxiliary model output.

scope controls the particle axis and MPI reduction. extensive records whether a configuration value scales with system size so an engine can normalize it correctly. Extensivity does not change reduction scope and is not valid for PARTICLE outputs.

extensive: bool

Alias for field number 2

name: str

Alias for field number 0

scope: OutputScope

Alias for field number 1

class Exporter

Exports a potential model to a serialized StableHLO bundle.

Subclasses select a graph representation and implement energy_fn(). The deployment documentation contains complete export examples.

Variables:

• graph_type (chemtrain.deploy.graphs.NeighborList) – Graph representation used for engine neighbor data.

• nbr_order (List[int]) – Two neighbor orders used with Newton on and Newton off, respectively.

• r_cutoff (float) – Model cutoff radius.

• unit_style (str) – Unit style used for positions and energies. Force units follow from the length and energy units.

• position_dtype (Any) – Floating-point dtype of the position input and its force and virial derivatives. The energy function determines the dtypes of its calculations and returned values. Engine adapters may use a different floating-point dtype because the connector converts at the engine ABI boundary.

• has_aux (bool) – If True, energy_fn returns (particle_energy, outputs).

• particle_fields (Tuple[chemtrain.deploy.exporter.ParticleField, ...]) – Ordered descriptors for additional scalar int32 arrays supplied through particle_data. Zero-based species is always provided and must not be registered explicitly.

• global_fields (Tuple[chemtrain.deploy.exporter.GlobalField, ...]) – Ordered descriptors for scalar inputs supplied through global_data. Global field values may vary between force calls.

• output_fields (Tuple[chemtrain.deploy.exporter.OutputField, ...]) – Ordered descriptors for auxiliary outputs and their PARTICLE, LOCAL, or GLOBAL scope. Configuration outputs also declare whether they are extensive. Shapes are inferred during tracing.

• communication_required (bool) – If True, callers must select the communication-enabled variant.

• include_pair_type (bool) – If True, exposes engine topology categories as graph.pair_type.

abstractmethod energy_fn(position, particle_data, graph, comm=None)

Computes particle energies for positions and a graph.

Parameters:

• position – (N, dim) Array of particle positions, including ghost atoms that are not within the local domain.

• particle_data – Mapping containing the built-in zero-based scalar int32 species array and every registered additional scalar array, with one value per particle.

• graph – Graph representation of the neighborhood around atoms. If include_pair_type is enabled, graph.pair_type contains aligned topology categories, with zero used for unclassified and padding entries.

• comm – Optional communication interface. Models that use communication call comm.gather or comm.reduce at fixed locations in the traced computation.

Returns:

An energy contribution for each particle. If has_aux is True, return (energy, outputs). Output keys must match output_fields. PARTICLE outputs start with the particle axis, while LOCAL and GLOBAL outputs describe one configuration.

export(*, communication=False, custom_calls=(), platforms=('cuda',))

Exports model variants specialized by communication and Newton mode.

Every bundle contains the two variants without model communication, one for each Newton mode. communication=True also adds the communication-enabled Newton-on variant. custom_calls lists additional model-specific FFI targets, while platforms selects CPU, CUDA, or both implementations in each variant.

Return type:

None

graph_type

alias of SimpleSparseNeighborList

position_dtype

alias of float32

save(file)

Saves the exported protocol buffer to file.

Return type:

None

Custom Calls

The following targets support exported models that use OpenEquivariance:

OPENEQUIVARIANCE_CUSTOM_CALLS = ('conv_forward', 'conv_backward', 'conv_double_backward')

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable’s items.

If the argument is a tuple, the return value is the same object.