Inference Context

class mlip.models.inference_context.InferenceContext(dataset_name: str | None = None, dataset_idx: int | None = None, charge: int | None = None, spin_multiplicity: int | None = None)

Graph-level metadata applied to a graph before inference.

Used to populate routing globals (charge, spin multiplicity, dataset index) that MoE-enabled models need for expert selection or multi-head models.

dataset_name

Human-readable dataset identifier. Must match an entry in DatasetInfo.dataset_name. Resolved to a dataset_idx via resolve().

Type:

str | None

dataset_idx

Numeric index of the dataset. Equivalent to dataset_name; if both are given they must agree. resolve() fills whichever is missing from the other.

Type:

int | None

charge

Total system charge to tag on each graph’s globals.charge. Only applied if set.

Type:

int | None

spin_multiplicity

Spin multiplicity to tag on each graph’s globals.spin_multiplicity. Only applied if set.

Type:

int | None

resolve(dataset_info: DatasetInfo) → InferenceContext

Return a new InferenceContext with dataset_name and dataset_idx cross-filled against the model’s known datasets.

Raises ValueError when the name is unknown, the index is out-of-range, or the two disagree. Single-dataset models (where dataset_info.dataset_name is None) are pass-through.

mlip.models.inference_context.apply_inference_context_to_graph(graph: Graph, inference_context: InferenceContext) → Graph

Populate graph globals from an inference context.

Returns a new graph with the relevant global fields (dataset_idx, charge, spin_multiplicity) broadcast to [n_graphs].