API Reference #

add_eval_phase(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, callbacks: list[Callback] | None = None) → PhaseGroup[source]#

Constructs and registers a new evaluation phase.

Parameters:

label (str) – A label to assign to this evaluation phase. Must be unique to already registered phases.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – An optional list of losses to be applied during evaluation.
active_nodes (list[GraphNode] | None, optional) – A list of active GraphNodes in this phase of the experiment. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

Returns:

self

Return type:

add_evaluation(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, callbacks: list[Callback] | None = None) → PhaseGroup#

Constructs and registers a new evaluation phase.

Parameters:

label (str) – A label to assign to this evaluation phase. Must be unique to already registered phases.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – An optional list of losses to be applied during evaluation.
active_nodes (list[GraphNode] | None, optional) – A list of active GraphNodes in this phase of the experiment. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

Returns:

self

Return type:

add_group(group: PhaseGroup) → PhaseGroup[source]#

Register a sub-group to this group.

Description:: Registers a new phase to be run wihtin this phase group. Phases are executed in the order they are added.

Parameters:: group (PhaseGroup) – A fully constructed phase group to nest under this one.
Returns:: self
Return type:: PhaseGroup

add_item(item: ExperimentPhase | PhaseGroup)[source]#

Registers a phase or group to this collection.

Parameters:: item (ExperimentPhase | PhaseGroup) – New phase or phase group to append. The item must have a unique label relative to its parent group.

add_items(items: Iterable[ExperimentPhase | PhaseGroup])[source]#

Register phases or groups to this collection.

Parameters:: items (Iterable[ExperimentPhase | PhaseGroup]) – New phases or groups to append. All items must have a unique label relative to its parent group.

add_phase(phase: ExperimentPhase) → PhaseGroup[source]#

Register a phase to this group.

Description:: Registers a new phase to be run wihtin this phase group. Phases are executed in the order they are added.

Parameters:: phase (ExperimentPhase) – A fully constructed ExperimentPhase.
Returns:: self
Return type:: PhaseGroup

add_train(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None) → PhaseGroup#

Constructs and registers a new training phase.

Parameters:

label (str) – A label to assign to this training phase. Must be unique to already registered phases.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – A list of losses to be applied during this training phase.
n_epochs (int) – Number of epochs to perform.
active_nodes (list[str | GraphNode] | None, optional) – A list of active GraphNodes in this phase of the experiment. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_schedule (str | BatchSchedulingPolicy, optional) –
Defines how batches from multiple samplers are scheduled during training. This is only relevant if more than one sampler is defined in input_sources.

Let samplers S1 and S2 produce: S1 = [b1, b2, b3] and S2 = [c1, c2]

The outputs of each policy is given below:
- ”zip_strict”: (b1, c1), (b2, c2)
- ”zip_cycle”: (b1, c1), (b2, c2), (b3, c1)
- ”alternate_strict”: b1, c1, b2, c2
- ”alternate_cycle”: b1, c1, b2, c2, b3, c1
See also BatchSchedulingPolicy.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

Returns:

self

Return type:

add_train_phase(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None) → PhaseGroup[source]#

Constructs and registers a new training phase.

Parameters:

label (str) – A label to assign to this training phase. Must be unique to already registered phases.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – A list of losses to be applied during this training phase.
n_epochs (int) – Number of epochs to perform.
active_nodes (list[str | GraphNode] | None, optional) – A list of active GraphNodes in this phase of the experiment. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_schedule (str | BatchSchedulingPolicy, optional) –
Defines how batches from multiple samplers are scheduled during training. This is only relevant if more than one sampler is defined in input_sources.

Let samplers S1 and S2 produce: S1 = [b1, b2, b3] and S2 = [c1, c2]

The outputs of each policy is given below:
- ”zip_strict”: (b1, c1), (b2, c2)
- ”zip_cycle”: (b1, c1), (b2, c2), (b3, c1)
- ”alternate_strict”: b1, c1, b2, c2
- ”alternate_cycle”: b1, c1, b2, c2, b3, c1
See also BatchSchedulingPolicy.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

Returns:

self

Return type:

add_training(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None) → PhaseGroup#

Constructs and registers a new training phase.

Parameters:

label (str) – A label to assign to this training phase. Must be unique to already registered phases.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – A list of losses to be applied during this training phase.
n_epochs (int) – Number of epochs to perform.
active_nodes (list[str | GraphNode] | None, optional) – A list of active GraphNodes in this phase of the experiment. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_schedule (str | BatchSchedulingPolicy, optional) –
Defines how batches from multiple samplers are scheduled during training. This is only relevant if more than one sampler is defined in input_sources.

Let samplers S1 and S2 produce: S1 = [b1, b2, b3] and S2 = [c1, c2]

The outputs of each policy is given below:
- ”zip_strict”: (b1, c1), (b2, c2)
- ”zip_cycle”: (b1, c1), (b2, c2), (b3, c1)
- ”alternate_strict”: b1, c1, b2, c2
- ”alternate_cycle”: b1, c1, b2, c2, b3, c1
See also BatchSchedulingPolicy.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

Returns:

self

Return type:

property all: list[ExperimentPhase | PhaseGroup]#

All registered phases and phase groups in execution order.

This returns the top-level structure (i.e., include phase groups). Use .flatten() to unravel all nested groups into a single list of phases, in execution order.

clear() → PhaseGroup[source]#

Removes all phases and groups registered within this group.

Returns:: self
Return type:: PhaseGroup

flatten() → list[ExperimentPhase][source]#

Flattens all registered phase groups.

Description:: All phase groups are flattened into a their underlying phases. The returned list represents the execution order of all individual experiment phases.

Returns:: All flattened phased in execution order.
Return type:: list[ExperimentPhase]

classmethod from_config(config: dict[str, Any]) → PhaseGroup[source]#

Construct a phase group from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration details. Keys must be strings.
Returns:: Reconstructed phase group.
Return type:: PhaseGroup

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this phase group.

Returns:: Configuration used to reconstruct the phase group.
Return type:: dict[str, Any]

get_eval_phase(key: int | str) → EvalPhase[source]#

Retrieves a registered evaluation phase from this group.

Parameters:: key (str, int) – The label or registered position of the phase to return.
Returns:: The specified phase.
Return type:: EvalPhase

get_group(key: int | str) → PhaseGroup[source]#

Retrieves a registered sub-roup from this group.

Parameters:: key (str, int) – The label or registered position of the group to return.
Returns:: The specified sub-group.
Return type:: PhaseGroup

get_phase(key: int | str) → ExperimentPhase[source]#

Retrieves a registered phase from this group.

Parameters:: key (str, int) – The label or registered position of the phase to return.
Returns:: The specified phase.
Return type:: ExperimentPhase

get_train_phase(key: int | str) → TrainPhase[source]#

Retrieves a registered training phase from this group.

Parameters:: key (str, int) – The label or registered position of the phase to return.
Returns:: The specified phase.
Return type:: TrainPhase

items() → Iterable[tuple[str, ExperimentPhase | PhaseGroup]][source]#

A generator over label-phase/group pairs.

The returned items are keyed by their unique labels. They are not returned in execution order.

Yields:: tuple[str, ExperimentPhase | PhaseGroup] – Tuple of phase/group and its label.

property phase_groups: dict[str, PhaseGroup]#

Only the top-level registered phase groups, keyed by their labels.

The returned dict does not encode execution order. Use all() to get all phases and phase groups as executed.

property phases: dict[str, ExperimentPhase]#

Only the top-level registered phases, keyed by their labels.

The returned dict does not encode execution order. Use all() to get all phases and phase groups as executed.

remove_group(group: int | str | PhaseGroup) → PhaseGroup[source]#

Remove a sub-group from this group.

Description:: Removes a registered group by index, label, or instance.

Parameters:: group (str | int | PhaseGroup) – Phase position, label, or instance to remove.
Returns:: self
Return type:: PhaseGroup

remove_item(item: int | str | ExperimentPhase | PhaseGroup)[source]#

Remove a phase or phase group from this collection.

Description:: Removes a registered phase by its positional index, its label, or its instance.

Parameters:: item (int | str | ExperimentPhase | PhaseGroup) – Item position, label, or instance to remove.

remove_items(items: list[int | str | ExperimentPhase | PhaseGroup])[source]#

Remove phases or phase groups from this collection.

Description:: Removes all specified items from this collection.

Parameters:: items (list[int | str | ExperimentPhase | PhaseGroup]) – List of phase/group position, label, or instance to remove.

remove_phase(phase: int | str | ExperimentPhase) → PhaseGroup[source]#

Remove a phase from this group.

Description:: Removes a registered phase by index, label, or instance.

Parameters:: phase (str | int | ExperimentPhase) – Phase position, label, or instance to remove.
Returns:: self
Return type:: PhaseGroup

class PhaseGroupResults(label: str, _results: dict[str, ~modularml.core.experiment.results.phase_results.PhaseResults | ~modularml.core.experiment.results.group_results.PhaseGroupResults] = <factory>)[source]#

Bases: object

Hierarchical results container matching the structure of a PhaseGroup.

Description:

PhaseGroupResults stores PhaseResults and nested PhaseGroupResults in the same hierarchy as the PhaseGroup that produced them. Results are stored in insertion order matching execution order.

Provides convenience methods for:

Accessing results by phase or group label
Flattening nested results into a single-level mapping
Iterating over results in execution order

label

Phase-group label associated with this result tree.

Type:: str

_results

Ordered mapping of labels to phase or nested group results.

Type:: dict[str, PhaseResults | PhaseGroupResults]

Example

Accessing phase group results

>>> # Access results by label
>>> train_res = group_results.get_phase_result("train")  
>>> eval_res = group_results.get_phase_result("eval")  
>>> # Flatten nested structure
>>> flat = group_results.flatten()  
>>> for label, phase_res in flat.items():  
...     print(f"{label}: {phase_res!r}")
>>> # Iterate in execution order
>>> for label, result in group_results.items():  
...     print(label, type(result).__name__)

__init__(label: str, _results: dict[str, ~modularml.core.experiment.results.phase_results.PhaseResults | ~modularml.core.experiment.results.group_results.PhaseGroupResults] = <factory>) → None#

add_result(result: PhaseResults | PhaseGroupResults)[source]#

Record a phase or group result.

Parameters:

result (PhaseResults | PhaseGroupResults) – The result to add. Must have a unique label within this group.

Raises:

TypeError – If result is not a PhaseResults or PhaseGroupResults.
ValueError – If a result with the same label already exists.

flatten() → dict[str, PhaseResults][source]#

Flatten all nested groups into a single-level dict.

Description:

Recursively unravels the hierarchy of PhaseGroupResults into a flat mapping of phase labels to their PhaseResults. The returned dict preserves execution order.

All phase labels must be unique across the entire hierarchy. If duplicate labels are found, a ValueError is raised.

Returns:: A flat mapping of phase labels to results in execution order.
Return type:: dict[str, PhaseResults]
Raises:: ValueError – If duplicate phase labels exist across the hierarchy.

Example

All internal PhaseResults can be flattened to a single list:

>>> flat = group_results.flatten()  
>>> for label, phase_res in flat.items():  
...     print(f"{label}: {type(phase_res).__name__}")

get_eval_result(label: str | None = None) → EvalResults[source]#

Retrieve a EvalResults by its label.

Parameters:

label (str) – The evaluation phase label to look up. Auto-detected if omitted. Defaults to None.

Returns:

The results for the specified phase.

Return type:

EvalResults

Raises:

KeyError – If no result exists with the given label.
TypeError – If the result is not of type EvalResults.

get_group_result(label: str | None = None) → PhaseGroupResults[source]#

Retrieve a nested PhaseGroupResults by its label.

Parameters:

label (str) – The group label to look up. Auto-detected if omitted. Defaults to None.

Returns:

The results for the specified group.

Return type:

PhaseGroupResults

Raises:

KeyError – If no result exists with the given label.
TypeError – If the result is a PhaseResults, not PhaseGroupResults.

get_phase_result(label: str) → PhaseResults[source]#

Retrieve a PhaseResults by its label.

Parameters:

label (str) – The phase label to look up.

Returns:

The results for the specified phase.

Return type:

PhaseResults

Raises:

KeyError – If no result exists with the given label.
TypeError – If the result is not of type PhaseResults.

get_train_result(label: str | None = None) → TrainResults[source]#

Retrieve a TrainResults by its label.

Parameters:

label (str, optional) – The training phase label to look up. Auto-detected if omitted. Defaults to None.

Returns:

The results for the specified phase.

Return type:

TrainResults

Raises:

KeyError – If no result exists with the given label.
TypeError – If the result is not of type TrainResults.

property group_results: dict[str, PhaseGroupResults]#

Only the top-level PhaseGroupResults entries, keyed by label.

Description:: Returns only the nested PhaseGroupResults (not PhaseResults) at this level of the hierarchy. The returned dict does not encode execution order.

items() → Iterator[tuple[str, PhaseResults | PhaseGroupResults]][source]#

Iterate over label-result pairs in execution order.

Returns:: Iterator over label/result pairs.
Return type:: Iterator[tuple[str, PhaseResults | PhaseGroupResults]]

label: str#

property labels: list[str]#

All top-level result labels in insertion order.

Returns:: Ordered labels.
Return type:: list[str]

property phase_results: dict[str, PhaseResults]#

Only the top-level PhaseResults entries, keyed by label.

Description:: Returns only the PhaseResults (not nested PhaseGroupResults) at this level of the hierarchy. The returned dict does not encode execution order.

class Checkpointing(*, mode: Literal['memory', 'disk'] = 'memory', save_on: str | list[str] = 'epoch_end', every_n: int = 1, directory: Path | str | None = None, name_template: str | None = None, max_to_keep: int | None = None, overwrite: bool = True)[source]#

Bases: object

Automatic state checkpointing configuration and tracker.

Description:

Manages when/where to save and tracks saved checkpoints. The actual save/restore operations are performed by the parent object (TrainPhase for model state, Experiment for full experiment state).

When attached to a TrainPhase, model graph state is saved at training lifecycle hooks (epoch boundaries, etc.). When attached to an Experiment, the full experiment (including results) is saved at execution lifecycle hooks (phase/group boundaries).

mode

Storage mode (memory or disk).

Type:: str

save_on

Lifecycle hooks that trigger saving.

Type:: list[str]

every_n

Frequency multiplier applied to hook counts.

Type:: int

directory

Destination directory for disk checkpoints.

Type:: Path | None

name_template

Template used to format checkpoint names.

Type:: str | None

max_to_keep

Limit on retained checkpoints (None keeps all).

Type:: int | None

overwrite

Whether disk checkpoints may overwrite existing files.

Type:: bool

__init__(*, mode: Literal['memory', 'disk'] = 'memory', save_on: str | list[str] = 'epoch_end', every_n: int = 1, directory: Path | str | None = None, name_template: str | None = None, max_to_keep: int | None = None, overwrite: bool = True) → None[source]#

Initialize a Checkpointing configuration.

Parameters:

mode (Literal["memory", "disk"], optional) – Storage mode for checkpoints. memory uses in-memory state snapshots (fast, no I/O). disk writes checkpoint files to the filesystem (persistent). Defaults to memory.
save_on (str | list[str], optional) –
Lifecycle hook(s) at which to save checkpoints. Valid values depend on where the Checkpointing is attached:
- TrainPhase: phase_start, phase_end, epoch_start, epoch_end, batch_start, batch_end
- Experiment: phase_start, phase_end, group_start, group_end
Defaults to epoch_end.
every_n (int, optional) – Save a checkpoint every every_n occurrences of the hook. For epoch_end, this means every N epochs. Defaults to 1.
directory (Path | str | None, optional) – Directory for disk-mode checkpoints. Ignored in memory mode. When attached to a TrainPhase and left as None, the directory is inherited from the parent Experiment at execution time (saved under a {phase.label}/ subdirectory).
name_template (str | None, optional) – Template for checkpoint file names. Available placeholders depend on context. For TrainPhase: {phase}, {epoch}, {batch}. For Experiment: {label}. If None, a context-appropriate default is assigned by the parent’s set_checkpointing() method. Defaults to None.
max_to_keep (int | None, optional) – Maximum number of checkpoints to retain. When exceeded, the oldest checkpoint is removed. If None, all checkpoints are kept. Defaults to None.
overwrite (bool, optional) – Whether to overwrite existing checkpoint files in disk mode. Defaults to True.

property directory: Path | None#

Get the checkpoint directory.

Returns:: Disk directory used when mode is disk.
Return type:: Path | None

property every_n: int#

Get the save frequency.

Returns:: Number of hook invocations between saves.
Return type:: int

format_name(**kwargs: Any) → str[source]#

Format the name template with provided key-value pairs.

Parameters:: **kwargs (Any) – Values to substitute into the name template. Common keys: label, key, phase, epoch, batch.
Returns:: The formatted checkpoint name.
Return type:: str

classmethod from_config(config: dict) → Checkpointing[source]#

Construct from config data.

Parameters:: config (dict) – Serialized configuration produced by get_config().
Returns:: Reconstructed checkpointing helper.
Return type:: Checkpointing

get_config() → dict[str, Any][source]#: Return configuration details required to reconstruct this object.

get_path(key: int | str) → Path | None[source]#

Return the disk checkpoint path for a given key.

Only available in disk mode. Returns None if no path exists or if in memory mode.

Parameters:: key (int | str) – Identifier previously supplied to record_disk().
Returns:: Checkpoint path if it exists on disk.
Return type:: Path | None

get_state(key: int | str) → dict[str, Any] | None[source]#

Return the in-memory state dict for a given key.

Only available in memory mode. Returns None if no state exists or if in disk mode.

Parameters:: key (int | str) – Identifier previously supplied to record_memory().
Returns:: Stored state if available.
Return type:: dict[str, Any] | None

has_key(key: int | str) → bool[source]#

Check whether a checkpoint exists for the given key.

Parameters:: key (int | str) – Identifier to query.
Returns:: True if the checkpoint is available.
Return type:: bool

property mode: str#

Get the storage mode.

Returns:: Either memory or disk.
Return type:: str

property name_template: str | None#

Get the template for checkpoint names.

Returns:: Template string with placeholders such as {phase}.
Return type:: str | None

property overwrite: bool#

Determine if disk checkpoints overwrite existing files.

Returns:: True if overwrites are allowed.
Return type:: bool

record_disk(key: int | str, path: Path) → None[source]#

Record a disk checkpoint path.

Parameters:

key (int | str) – Identifier for this checkpoint.
path (Path) – The saved checkpoint file path.

record_memory(key: int | str, state: dict[str, Any]) → None[source]#

Record an in-memory state snapshot.

Parameters:

key (int | str) – Identifier for this checkpoint (e.g. epoch index or phase label).
state (dict[str, Any]) – The state dictionary to store.

reset() → None[source]#

Clear all tracked states and counters.

Called by the parent at the start of a new phase or experiment run to ensure fresh checkpointing state.

property save_on: list[str]#

Get the hook names at which checkpoints are saved.

Returns:: Lifecycle hooks configured for saving.
Return type:: list[str]

property saved_keys: list[int | str]#

List cache keys for which checkpoints exist.

Returns:: Keys in sorted order.
Return type:: list[int | str]

should_save(hook: str) → bool[source]#

Check if a save should occur for this hook.

Checks whether hook is in save_on and whether the every_n condition is satisfied.

Parameters:: hook (str) – The lifecycle hook name.
Returns:: True if a checkpoint should be saved.
Return type:: bool

static validate_placeholders(template: str, allowed: frozenset[str], context_name: str) → None[source]#

Validate that a name template only uses allowed placeholders.

Parameters:

template (str) – The name template string to validate.
allowed (frozenset[str]) – Set of allowed placeholder names.
context_name (str) – Name of the context (for error messages).

Raises:

ValueError – If the template contains disallowed placeholders.

class InputBinding(node_id: str, upstream_ref: FeatureSetReference, split: str | None = None, sampler: BaseSampler | None = None, stream: str = 'default')[source]#

Bases: object

A phase-specific binding of input data to a head GraphNode.

Description:: An InputBinding exists within a single Experiment phase. It defines an attachment of a sampler (or direct pass-through) to an existing graph edge between a FeatureSet and a head GraphNode, optionally restricted to a FeatureSet split.

node_id

ID of the head GraphNode on which we are defining input for.

Type:: str

upstream_ref

Which upstream reference of the head node this binding applies to.

Type:: FeatureSetReference

split

If defined, only data from this split is used.

Type:: str, optional

sampler

A sampler to use in feeding the source data to the head node.

Type:: BaseSampler, optional

stream

If a sampler with multiple output streams is used, this defines the exact stream of data to feed to the head node.

Type:: str, optional

__init__(node_id: str, upstream_ref: FeatureSetReference, split: str | None = None, sampler: BaseSampler | None = None, stream: str = 'default') → None#

Create an InputBinding for an evaluation phase.

Description:

This method creates a phase-specific binding that directly feeds data from an upstream FeatureSet into a head GraphNode without applying a sampler.

Evaluation bindings typically iterate over all samples in a split (or the full FeatureSet) and are used for validation, testing, or inference.

Parameters:

node (GraphNode | str) – The head GraphNode that will receive input data during evaluation. Accepted values: - GraphNode instance - GraphNode label (str) - GraphNode ID (str)
upstream (FeatureSet | FeatureSetView | str | None) –
Identifies which upstream FeatureSet connection of node this binding applies to. Accepted values: - FeatureSet instance - FeatureSetView instance - FeatureSet node ID or label (str) - None, only if node has exactly one upstream FeatureSet

If the node has multiple upstream FeatureSets, this argument is required to disambiguate which input is being bound.
split (str, optional) – Optional split name of the upstream FeatureSet (e.g. “val”, “test”). If provided, only rows from this split are used. If None, the entire FeatureSet is evaluated.

Returns:

A fully specified evaluation InputBinding that can be passed directly to an EvalPhase.

Return type:

InputBinding

Create an InputBinding for a training phase.

Description:

This method creates a phase-specific binding that attaches a sampler between an upstream FeatureSet and a head GraphNode.

Conceptually, this binding modifies an existing graph edge (FeatureSet -> GraphNode) by inserting a sampler that controls how samples are batched and fed into the node during training.

No data is materialized at construction time. The sampler is only executed when the training phase runs.

Parameters:

node (GraphNode | str) – The head GraphNode that will receive input data during training. Accepted values: - GraphNode instance - GraphNode label (str) - GraphNode ID (str)
sampler (BaseSampler) – The sampler used to generate batches from the upstream FeatureSet (e.g., random batches, contrastive roles, paired samples).
upstream (FeatureSet | FeatureSetView | str | None) –
Identifies which upstream FeatureSet connection of node this binding applies to. Accepted values: - FeatureSet instance - FeatureSetView instance - FeatureSet node ID or label (str) - None, only if node has exactly one upstream FeatureSet

If the node has multiple upstream FeatureSets, this argument is required to disambiguate which input is being bound.
split (str, optional) – Optional split name of the upstream FeatureSet (e.g. “train”, “val”). If provided, only rows from this split are sampled. If None, the entire FeatureSet is used.
stream (str, optional) – Output stream name from the sampler to feed into node. Required only if the sampler produces multiple streams. Defaults to STREAM_DEFAULT.

Returns:

A fully specified training InputBinding that can be passed directly to a TrainPhase.

Return type:

InputBinding

classmethod from_config(config: dict) → ExperimentPhase[source]#

Construct a phase from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration details. Keys must be strings.
Returns:: Reconstructed phase.
Return type:: ExperimentPhase

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this binding.

Returns:: Configuration used to reconstruct the binding.
Return type:: dict[str, Any]

materialize_batches(*, show_progress: bool = True) → list[BatchView][source]#

Executes sampling of the source data defined by this binding.

Parameters:: show_progress (bool, optional) – Whether to show a progress bar of the batch construction process.
Returns:: The materialized batches for the sampler and stream defined by this binding.
Return type:: list[BatchView]

node_id: str#

resolve_input_view() → FeatureSetView[source]#

Resolves the FeatureSetView for the upstream_ref.

Returns:: A view of the FeatureSet specified by upstream_ref. If split is defined, the returned view is restricted to only the indices of the split.
Return type:: FeatureSetView

sampler: BaseSampler | None = None#

split: str | None = None#

stream: str = 'default'#

upstream_ref: FeatureSetReference#

class ResultRecording(value)[source]#

Bases: Enum

Controls which execution contexts are retained in TrainResults.

ALL:: Record every batch of every epoch (default). Provides full access to all outputs, losses, and tensors across the entire training run but uses more memory.
LAST:: Keep only the final epoch’s execution contexts. When an EarlyStopping callback with restore_best=True is active, “last” is interpreted as the best epoch.
NONE:: Do not record execution contexts at all. Scalar metrics (e.g. train_loss, val_loss) are still logged to the MetricStore and remain accessible.

ALL = 'all'#

LAST = 'last'#

NONE = 'none'#

classmethod from_value(value: str | ResultRecording) → ResultRecording[source]#

Normalize strings or enums to a ResultRecording.

Parameters:: value (str | ResultRecording) – Input value to normalize.
Returns:: Canonical enum member.
Return type:: ResultRecording
Raises:: ValueError – If value cannot be mapped to a valid member.

class FitPhase(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, freeze_after_fit: bool = True, callbacks: list[Callback] | None = None)[source]#

Bases: ExperimentPhase

Phase that fits batch-fit model nodes on the complete dataset.

Description:

FitPhase is designed for scikit-learn models (and similar) that require all training data at once via .fit(X, y) rather than iterative mini-batch gradient updates.

Unlike TrainPhase, FitPhase has no epochs or sampling. It yields a single ExecutionContext containing the entire dataset from the specified split(s).

By default, fitted nodes are frozen after fitting so that downstream gradient-trained nodes can use their outputs without interference.

__init__(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, freeze_after_fit: bool = True, callbacks: list[Callback] | None = None)[source]#

Initialize a new fit phase for the experiment.

Notes

All input_sources must originate from the same upstream FeatureSet. If multiple FeatureSets need to be fitted, they must be done so in separate FitPhases.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph. All bindings must resolve to the same FeatureSet.
losses (list[AppliedLoss], optional) – A list of losses to compute after fitting (for metrics only).
active_nodes (list[GraphNode] | None, optional) – A list of GraphNodes to fit. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
freeze_after_fit (bool, optional) – Whether to freeze fitted nodes after .fit() completes. Defaults to True.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

classmethod from_config(config: dict) → FitPhase[source]#

Construct a FitPhase from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration details. Keys must be strings.
Returns:: Reconstructed phase.
Return type:: FitPhase

classmethod from_split(label: str, *, split: str, losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, freeze_after_fit: bool = True, callbacks: list[Callback] | None = None) → FitPhase[source]#

Initialize a new fit phase for a given FeatureSet split.

Notes

All active head nodes must input from the defined split.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
split (str) – The FeatureSet split name to fit on (e.g., “train”).
losses (list[AppliedLoss], optional) – A list of losses to compute after fitting (for metrics only).
active_nodes (list[GraphNode] | None, optional) – A list of GraphNodes to fit. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
freeze_after_fit (bool, optional) – Whether to freeze fitted nodes after .fit() completes. Defaults to True.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this phase.

Returns:: Configuration used to reconstruct the phase.
Return type:: dict[str, Any]

iter_execution(*, results: FitResults | None = None) → Iterator[ExecutionContext][source]#

Iterate over execution steps for this fit phase.

Description:: Generates a single ExecutionContext containing the entire dataset from the specified split. No batching or epochs are used.

Parameters:: results (FitResults | None, optional) – Optional container in which results will be registered.
Yields:: ExecutionContext – A single execution context containing all data, suitable for ModelGraph.fit_step(ctx).

class FitResults(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>)[source]#

Bases: PhaseResults

Results container for a single FitPhase execution.

Description:

FitResults wraps the outputs of a FitPhase, which executes a single pass over the complete dataset (epoch=0, batch=0). This class provides convenience methods for:

Loss aggregation from the fit pass
Access to fitted model outputs

label

Phase label inherited from PhaseResults.

Type:: str

__init__(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>) → None#

aggregated_losses(node: str | GraphNode, *, reducer: Literal['mean', 'sum'] = 'mean') → dict[str, float][source]#

Retrieve losses from the fit phase.

Parameters:

node (str | GraphNode) – The node to filter losses to.
reducer (Literal['mean', 'sum']) – How losses should be aggregated. Defaults to “mean”.

Returns:

Losses keyed by the AppliedLoss label.

Return type:

dict[str, float]

source_view(node: str | GraphNode, *, role: str = 'default') → FeatureSetView[source]#

Get the single source FeatureSetView for the given node.

Description:

Convenience method for the common case where a node has exactly one upstream FeatureSet. Raises ValueError if multiple upstream FeatureSets exist.

Note that the returned views contain only unique sample UUIDs used in generating these phase results. They are not a 1-to-1 mapping of result sample to source sample. Use tensors() to get exact execution data.

Parameters:

node (str | GraphNode) – The node to trace upstream from.
role (str, optional) – Restrict to samples from this role only. Defaults to “default”.
batch (int | None, optional) – Restrict to samples from this batch only.

Returns:

A view of the single upstream FeatureSet filtered to only the samples used during execution.

Return type:

FeatureSetView

Raises:

ValueError – If the node has multiple upstream FeatureSets.

source_views(node: str | GraphNode, *, role: str = 'default') → dict[str, FeatureSetView][source]#

Get the source FeatureSetViews that contributed data to the given node.

Description:

Traces the node back to its upstream FeatureSets, collects all unique sample UUIDs from execution results, and returns a view of each upstream FeatureSet filtered to only the samples used.

Note that the returned views contain only unique sample UUIDs used in generating these phase results. They are not a 1-to-1 mapping of result sample to source sample. Use tensors() to get exact execution data.

Parameters:

node (str | GraphNode) – The node to trace upstream from. Can be the node instance, its ID, or its label.
role (str, optional) – Restrict to samples from this role only. Defaults to “default”.
batch (int | None, optional) – Restrict to samples from this batch only.

Returns:

A mapping of FeatureSet label to FeatureSetView containing only the samples used during execution.

Return type:

dict[str, FeatureSetView]

stacked_batches(node: str | GraphNode, *, fmt: DataFormat | None = None) → Batch[source]#

Retrieve all batches for a node, concatenated into a single Batch.

Description:: Since a FitPhase executes all data under a single batch, this method simply returns the single executed batch.

Parameters:

node (str | GraphNode) – The node to retrieve batches for.
fmt (DataFormat | None, optional) – Format to cast tensor data to. Defaults to None.

Returns:

A single Batch containing concatenated data from all batches.

Return type:

Batch

Example

Accessing stacked Batch objects across all execution batches

>>> batch = fit_results.stacked_batches(  
...     node="output_node"
... )
>>> print(f"Total samples: {batch.batch_size}")  
>>> print(f"Outputs shape: {batch.outputs.shape}")  

stacked_tensors(node: str | GraphNode, domain: Literal['outputs', 'targets', 'tags', 'sample_uuids'], *, role: str = 'default', fmt: DataFormat | None = None, unscale: bool = False) → TensorLike[source]#

Retrieve tensors for a node, concatenated across all batches.

Description:: Since a FitPhase executes all data under a single batch, this method simply returns the specified domain of the single recorded execution context.

Parameters:

node (str | GraphNode) – The node to retrieve tensors for. Can be the node instance, its ID, or its label.
domain (Literal["outputs", "targets", "tags", "sample_uuids"]) – The domain of data to return: - outputs: the tensors produced by the node forward pass - targets: the expected output tensors (only for tail nodes) - tags: any tracked tags during the node’s forward pass - sample_uuids: the sample identifiers
role (str, optional) – If multi-role data, specifies which role to return. Defaults to “default”.
fmt (DataFormat | None, optional) – Format to cast returned tensors to. If None, uses as-produced format. Defaults to None.
unscale (bool, optional) – Whether to inverse any applied scalers. Only valid for tail nodes with domain in [“outputs”, “targets”]. Defaults to False.

Returns:

A single tensor containing concatenated data from all batches.

Return type:

TensorLike

Example

Getting concatenated tensors across all stacked batches:

>>> # Get all predictions stacked
>>> predictions = fit_results.stacked_tensors(  
...     node="output_node",
...     domain="outputs",
... )
>>> # Get targets, unscaled, as numpy
>>> targets = fit_results.stacked_tensors(  
...     node="output_node",
...     domain="targets",
...     fmt="np",
...     unscale=True,
... )

class TrainPhase(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[str | GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None, checkpointing: Checkpointing | None = None, result_recording: ResultRecording | str = ResultRecording.ALL)[source]#

Bases: ExperimentPhase

Phase that trains model graph nodes over one or more epochs.

__init__(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[str | GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None, checkpointing: Checkpointing | None = None, result_recording: ResultRecording | str = ResultRecording.ALL)[source]#

Initiallizes a new training phase for the experiment.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph.
losses (list[AppliedLoss]) – A list of losses to be applied during this training pahse.
n_epochs (int) – Number of epochs to perform.
active_nodes (list[str | GraphNode] | None, optional) – A list of GraphNodes to train in this training phase. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_schedule (str | BatchSchedulingPolicy, optional) –
Defines how batches from multiple samplers are scheduled during training. This is only relevant if more than one sampler is defined in input_sources.

Let samplers S1 and S2 produce: S1 = [b1, b2, b3] and S2 = [c1, c2]

The outputs of each policy is given below:
- ”zip_strict”: (b1, c1), (b2, c2)
- ”zip_cycle”: (b1, c1), (b2, c2), (b3, c1)
- ”alternate_strict”: b1, c1, b2, c2
- ”alternate_cycle”: b1, c1, b2, c2, b3, c1
See also BatchSchedulingPolicy.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.
checkpointing (Checkpointing | None, optional) – An optional Checkpointing callback that automatically saves model state at configurable lifecycle hook points. Unlike regular callbacks, this is configured as a phase-level argument rather than added manually via add_callback(). Defaults to None.
result_recording (ResultRecording | str, optional) – Controls which execution contexts are retained in the returned TrainResults. See ResultRecording for details. Defaults to ResultRecording.ALL.

add_callback(callback: Callback)[source]#

Add a callback to this training phase.

Parameters:: callback (Callback) – Callback to append.
Raises:: ValueError – If another callback of the same type and label exists.

property checkpointing: Checkpointing | None#: The Checkpointing instance configured for this phase, or None.

classmethod from_config(config: dict) → TrainPhase[source]#

Construct a phase from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration details. Keys must be strings.
Returns:: Reconstructed phase.
Return type:: ExperimentPhase

classmethod from_split(label: str, *, split: str, sampler: BaseSampler, losses: list[AppliedLoss], n_epochs: int = 1, active_nodes: list[str | GraphNode] | None = None, batch_schedule: BatchSchedulingPolicy | str = BatchSchedulingPolicy.ZIP_STRICT, callbacks: list[Callback] | None = None, checkpointing: Checkpointing | None = None, result_recording: ResultRecording | str = ResultRecording.ALL) → TrainPhase[source]#

Initiallizes a new training phase for a given FeatureSet split.

Notes

All active head nodes must input from the defined split. If the model graph has multiple head nodes that input from different FeatureSets, you will need to use the default TrainPhase constructor.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
split (str) – The FeatureSet split to train on.
sampler (BaseSampler, optional) – A sampler to use to generate batches from this split.
losses (list[AppliedLoss]) – A list of losses to be applied during this training pahse.
n_epochs (int) – Number of epochs to perform.
active_nodes (list[str | GraphNode] | None, optional) – A list of GraphNodes to train in this training phase. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_schedule (str | BatchSchedulingPolicy, optional) –
Defines how batches from multiple samplers are scheduled during training. This is only relevant if there is more than one head node.

Let samplers S1 and S2 produce: S1 = [b1, b2, b3] and S2 = [c1, c2]

The outputs of each policy is given below:
- ”zip_strict”: (b1, c1), (b2, c2)
- ”zip_cycle”: (b1, c1), (b2, c2), (b3, c1)
- ”alternate_strict”: b1, c1, b2, c2
- ”alternate_cycle”: b1, c1, b2, c2, b3, c1
See also BatchSchedulingPolicy.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.
checkpointing (Checkpointing | None, optional) – An optional Checkpointing callback that automatically saves model state at configurable lifecycle hook points. Defaults to None.
result_recording (ResultRecording | str, optional) – Controls which execution contexts are retained in the returned TrainResults. See ResultRecording for details. Defaults to ResultRecording.ALL.

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this phase.

Returns:: Configuration used to reconstruct the phase.
Return type:: dict[str, Any]

is_epoch_end() → bool[source]#: Whether current iter_execution state is at the end of an epoch.

iter_execution(*, results: TrainResults | None = None, show_sampler_progress: bool = True, show_training_progress: bool = True, persist_progress: bool = False, persist_epoch_progress: bool = False, val_loss_metric: str = 'val_loss') → Iterator[ExecutionContext][source]#

Iterate over all execution steps for this training phase.

This generator produces a sequence of ExecutionContext objects representing the full training schedule of the phase, across all epochs and all batch steps within each epoch.

The execution flow is:

Group input bindings into unique sampler executions via _build_sampler_executions(). Samplers with identical configuration, FeatureSet, and split are executed only once and shared across bindings.
For each sampler execution, obtain the number of materialized batches from its SampledView.
For each epoch:
1. Generate a step-wise batch schedule using _iter_schedule() according to self.batch_schedule.
2. For each schedule step, construct the inputs for all head-node bindings:
  - If a sampler is active in the current step, its corresponding batch is selected.
  - If a sampler is inactive (ALTERNATE policies), its inputs are replaced with a fully masked BatchView.
Yield an ExecutionContext containing:
- Phase label
- Epoch index
- Batch index (within the epoch)
- Resolved input BatchView objects for all bindings

Notes

ZIP scheduling policies always provide one batch per sampler per step.
ALTERNATE scheduling policies activate exactly one sampler per step; all others are masked.
No semantic alignment between samplers is performed here. Any required alignment (e.g., contrastive pairs or matched samples) must be handled inside the sampler itself via roles.
The yielded ExecutionContext objects are intended to be consumed directly by the ModelGraph training loop (e.g., ModelGraph.train_step(ctx)).

Parameters:

results (TrainResults | None, optional) – Optional container in which results will be registered.
show_sampler_progress (bool, optional) – Whether to show a progress bar for sampler batching. Defaults to True.
show_training_progress (bool, optional) – Whether to show a progress bar for training execution. Defaults to True.
persist_progress (bool, optional) – Whether to leave all progress bars visible after completion. Overrides all nested progress bar persistence settings. Defaults to IN_NOTEBOOK (True in notebooks, False in scripts).
persist_epoch_progress (bool, optional) – Whether to leave per-epoch training bars visible after completion. Defaults to IN_NOTEBOOK.
val_loss_metric (str, optional) – Name of a recorded validation loss metric to display in the progress bar. Results must be tracked and the metric must exist. If not, no validation loss field will be shown. Defaults to "val_loss".

Yields:

ExecutionContext – A fully specified execution context for a single batch step within a specific epoch of this training phase.

request_stop() → None[source]#

Request early termination of this training phase.

Description:: Sets an internal flag that is checked at the end of each epoch. When set, the training loop will break cleanly after the current epoch completes. Intended to be called by callbacks like EarlyStopping.

set_checkpointing(checkpointing: Checkpointing | None) → None[source]#

Attach or replace the Checkpointing configuration for this phase.

Validates that all save_on hooks are valid for a TrainPhase and that the name_template only uses allowed placeholders. If no name_template is set, the training default is applied.

Parameters:: checkpointing (Checkpointing | None) – The Checkpointing configuration, or None to disable.

class TrainResults(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>)[source]#

Bases: PhaseResults

Results container for a training phase.

Description:

TrainResults wraps the outputs of a TrainPhase, which executes multiple epochs with multiple batches per epoch. This class provides:

Access to training data keyed by epoch and batches
Direct access to validation losses/tensors from evaluation callbacks
Loss aggregation per epoch

Validation callbacks (kind=”evaluation”) are automatically detected and their results exposed through dedicated accessors.

label

Phase label.

Type:: str

_execution

Ordered execution contexts.

Type:: list[ExecutionContext]

_callbacks

Recorded callback outputs.

Type:: list[CallbackResult]

_metrics

Stored scalar metrics.

Type:: MetricStore

_series_cache

Cache of memoized AxisSeries queries.

Type:: dict[tuple, Any]

__init__(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>) → None#

property epoch_indices: list[int]#

Sorted list of recorded epoch indices.

Returns:: Epoch indices in ascending order.
Return type:: list[int]

property n_epochs: int#

The number of epochs executed during training.

Returns:: Total number of recorded epochs.
Return type:: int

class EvalPhase(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, batch_size: int | None = None, callbacks: list[Callback] | None = None)[source]#

Bases: ExperimentPhase

Phase that evaluates model outputs on a fixed FeatureSet split.

__init__(label: str, *, input_sources: list[InputBinding], losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, batch_size: int | None = None, callbacks: list[Callback] | None = None)[source]#

Initiallizes a new evaluation phase for the experiment.

Notes

All input_sources must originate from the same upstream FeatureSet. If multiple FeatureSets need to be evaluated, they must be done so in separate EvalPhases.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
input_sources (list[InputBinding]) – Input bindings for each head node in ModelGraph. All bindings must resolve to the same FeatureSet.
losses (list[AppliedLoss], optional) – A list of losses to be applied during this evaluation phase.
active_nodes (list[GraphNode] | None, optional) – A list of GraphNodes to run a forward phase on. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_size (int, optional) – If defined, limits the number of samples to using during a single forward pass. Otherwise, all samples are passed at once.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

classmethod from_config(config: dict) → EvalPhase[source]#

Construct a phase from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration details. Keys must be strings.
Returns:: Reconstructed phase.
Return type:: ExperimentPhase

classmethod from_split(label: str, *, split: str, losses: list[AppliedLoss] | None = None, active_nodes: list[GraphNode] | None = None, batch_size: int | None = None, callbacks: list[Callback] | None = None) → EvalPhase[source]#

Initiallizes a new evaluation phase for a given FeatureSet split.

Notes

All active head nodes must input from the defined split.

Parameters:

label (str) – A label to assign to this phase of the experiment. Used for logging.
split (str) – The FeatureSet split name to evaluate.
losses (list[AppliedLoss], optional) – A list of losses to be applied during this evaluation phase.
active_nodes (list[GraphNode] | None, optional) – A list of GraphNodes to run a forward phase on. Nodes can be listed via their ID, label, or with the actual node instance. If None, all nodes comprising the ModelGraph are used. Defaults to None.
batch_size (int, optional) – If defined, limits the number of samples to using during a single forward pass. Otherwise, all samples are passed at once.
callbacks (list[Callback] | None, optional) – An optional list of Callbacks to run during phase execution.

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this phase.

Returns:: Configuration used to reconstruct the phase.
Return type:: dict[str, Any]

iter_execution(*, results: EvalResults | None = None, show_eval_progress: bool = False, persist_progress: bool = False) → Iterator[ExecutionContext][source]#

Iterate over execution steps for this evaluation phase.

Description:: Generates ExecutionContext objects that cover the specified split of the upstream FeatureSet. The split is optionally chunked into batches based on batch_size to limit memory usage.

Parameters:

results (EvalResults | None, optional) – Optional container in which results will be registered.
show_eval_progress (bool, optional) – Whether to show a progress bar for eval batches. Defaults to False.
persist_progress (bool, optional) – Whether to leave all eval progress bars shown after they complete. Defaults to IN_NOTEBOOK (True if working in a notebook, False if in a Python script).

Yields:

ExecutionContext – Execution contexts suitable for forward-only execution via ModelGraph.eval_step(ctx).

class EvalResults(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>)[source]#

Bases: PhaseResults

Results container for a single forward-pass evaluation phase.

Description:

EvalResults wraps the outputs of an EvalPhase, which executes a single epoch (epoch=0) over multiple batches. This class provides convenience methods for:

Automatic tensor stacking across batches
Loss aggregation (sum/mean) over batches

All methods leverage the base PhaseResults query interface and use AxisSeries.collapse() for batch aggregation.

label

Phase label inherited from PhaseResults.

Type:: str

Example

Accessing EvalPhase results

>>> # Run evaluation
>>> eval_results = experiment.run_evaluation(phase=eval_phase)  
>>> # Get stacked outputs for a node (all batches concatenated)
>>> outputs = eval_results.stacked_tensors(  
...     node="output_node", domain="outputs"
... )
>>> # Get total loss across all batches
>>> total_loss = eval_results.aggregated_losses(  
...     node="output_node"
... )
>>> # Get all source data utilized in evaluation
>>> source_view = eval_results.source_view(node="output_node")  

__init__(label: str, _execution: list[ExecutionContext] = <factory>, _callbacks: list[CallbackResult] = <factory>, _metrics: MetricStore = <factory>) → None#

aggregated_losses(node: str | GraphNode, *, reducer: Literal['mean', 'sum'] = 'mean') → dict[str, float][source]#

Aggregates losses over all batches within this eval phase.

Parameters:

node (str | GraphNode) – The node to filter losses to. Can be the node instance, its ID, or its label.
reducer (Literal['mean', 'sum']) – How losses should be aggregated. Defaults to “mean”.

Returns:

Aggregated lossed, keyed by the AppliedLoss label.

Return type:

dict[str, float]

property batch_indices: list[int]#

Sorted list of recorded batch indices.

Returns:: Batch indices in ascending order.
Return type:: list[int]

property n_batches: int#

The number of batches executed during evaluation.

Returns:: Number of recorded batches.
Return type:: int

source_view(node: str | GraphNode, *, role: str = 'default', batch: int | None = None) → FeatureSetView[source]#

Get the single source FeatureSetView for the given node.

Description:

Convenience method for the common case where a node has exactly one upstream FeatureSet. Raises ValueError if multiple upstream FeatureSets exist.

Note that the returned views contain only unique sample UUIDs used in generating these phase results. They are not a 1-to-1 mapping of result sample to source sample. Use tensors() to get exact execution data.

Parameters:

node (str | GraphNode) – The node to trace upstream from.
role (str, optional) – Restrict to samples from this role only. Defaults to “default”.
batch (int | None, optional) – Restrict to samples from this batch only.

Returns:

A view of the single upstream FeatureSet filtered to only the samples used during execution.

Return type:

FeatureSetView

Raises:

ValueError – If the node has multiple upstream FeatureSets.

source_views(node: str | GraphNode, *, role: str = 'default', batch: int | None = None) → dict[str, FeatureSetView][source]#

Get the source FeatureSetViews that contributed data to the given node.

Description:

Traces the node back to its upstream FeatureSets, collects all unique sample UUIDs from execution results, and returns a view of each upstream FeatureSet filtered to only the samples used.

Note that the returned views contain only unique sample UUIDs used in generating these phase results. They are not a 1-to-1 mapping of result sample to source sample. Use tensors() to get exact execution data.

Parameters:

node (str | GraphNode) – The node to trace upstream from. Can be the node instance, its ID, or its label.
role (str, optional) – Restrict to samples from this role only. Defaults to “default”.
batch (int | None, optional) – Restrict to samples from this batch only.

Returns:

A mapping of FeatureSet label to FeatureSetView containing only the samples used during execution.

Return type:

dict[str, FeatureSetView]

stacked_batches(node: str | GraphNode, *, fmt: DataFormat | None = None) → Batch[source]#

Retrieve all batches for a node, concatenated into a single Batch.

Description:: Collects Batch objects from all evaluation batches and concatenates them using Batch.concat(). This provides access to all data domains (outputs, targets, tags, sample_uuids) plus role weights and masks in a single container.

Parameters:

node (str | GraphNode) – The node to retrieve batches for.
fmt (DataFormat | None, optional) – Format to cast tensor data to. Defaults to None.

Returns:

A single Batch containing concatenated data from all batches.

Return type:

Batch

Example

Getting concatenated batches across all execution batches

>>> batch = eval_results.stacked_batches(  
...     node="output_node"
... )
>>> print(f"Total samples: {batch.batch_size}")  
>>> print(f"Outputs shape: {batch.outputs.shape}")  

stacked_tensors(node: str | GraphNode, domain: Literal['outputs', 'targets', 'tags', 'sample_uuids'], *, role: str = 'default', fmt: DataFormat | None = None, unscale: bool = False) → TensorLike[source]#

Retrieve tensors for a node, concatenated across all batches.

Description:

Collects tensors from the specified domain across all evaluation batches and concatenates them along the batch dimension using backend-aware concatenation (torch.cat, np.concatenate, or tf.concat).

This is the primary method for retrieving complete evaluation outputs or targets in a single tensor.

Parameters:

node (str | GraphNode) – The node to retrieve tensors for. Can be the node instance, its ID, or its label.
domain (Literal["outputs", "targets", "tags", "sample_uuids"]) – The domain of data to return: - outputs: the tensors produced by the node forward pass - targets: the expected output tensors (only for tail nodes) - tags: any tracked tags during the node’s forward pass - sample_uuids: the sample identifiers
role (str, optional) – If multi-role data, specifies which role to return. Defaults to “default”.
fmt (DataFormat | None, optional) – Format to cast returned tensors to. If None, uses as-produced format. Defaults to None.
unscale (bool, optional) – Whether to inverse any applied scalers. Only valid for tail nodes with domain in [“outputs”, “targets”]. Defaults to False.

Returns:

A single tensor containing concatenated data from all batches.

Return type:

TensorLike

Example

Accessing all predictions across all eval batches:

>>> # Get all predictions stacked
>>> predictions = eval_results.stacked_tensors(  
...     node="output_node",
...     domain="outputs",
... )
>>> # Get targets, unscaled, as numpy
>>> targets = eval_results.stacked_tensors(  
...     node="output_node",
...     domain="targets",
...     fmt="np",
...     unscale=True,
... )

Execution Strategies#

class CrossValidation(*, label: str = 'CV', bindings: CVBinding | list[CVBinding], n_folds: int = 5, seed: int = 13, phase: TrainPhase | PhaseGroup | None = None, experiment: Experiment | None = None)[source]#

Bases: ExecutionStrategy

Cross-validation execution strategy.

Description:: Orchestrates repeated execution of an Experiment or PhaseGroup by remapping FeatureSet nodes to fold-specific train/validation splits defined via CVBinding objects.

__init__(*, label: str = 'CV', bindings: CVBinding | list[CVBinding], n_folds: int = 5, seed: int = 13, phase: TrainPhase | PhaseGroup | None = None, experiment: Experiment | None = None)[source]#

Initialize the cross-validation strategy.

Parameters:

label (str, optional) – Human-readable label applied to generated fold groups. Defaults to CV.
bindings (CVBinding | list[CVBinding]) – One or more CVBinding instances describing how each FeatureSet participates in folding.
n_folds (int, optional) – Number of folds to generate. Must be greater than or equal to 1. Defaults to 5.
seed (int, optional) – Random seed forwarded to RandomSplitter. Defaults to 13.
phase (TrainPhase | PhaseGroup | None, optional) – Optional TrainPhase or PhaseGroup template to run inside each fold. If omitted, the experiment execution plan is used.
experiment (Experiment | None, optional) – Experiment to execute. Defaults to the active experiment from ExperimentContext.

Raises:

TypeError – If no experiment is available or if phase has an invalid type.
ValueError – If n_folds or val_size settings are inconsistent.

run(*, show_fold_progress: bool = True, persist_progress: bool = False, **kwargs) → CVResults[source]#

Execute cross-validation across all folds.

Parameters:

show_fold_progress (bool, optional) – Whether to show a progress bar over fold execution. Defaults to True.
persist_progress (bool, optional) – Whether to keep progress bars visible after completion. Defaults to IN_NOTEBOOK (True in notebooks, False in scripts).
**kwargs – Additional display flags forwarded to Experiment.run_group().

Returns:

Cross-fold results container.

Return type:

CVResults

Bases: object

Configuration for cross-validation of a single FeatureSet.

Description:: Defines how a specified FeatureSet should participate in cross-validation. The head nodes of a ModelGraph are typically bound to a split or FeatureSet. This configuration maps fold-specific outputs back onto those bindings (for example, mapping train_split_name=’my_training’). Each fold then updates the expected split names with data belonging only to the current fold.

featureset

The FeatureSet to create folds from.

Type:: FeatureSet

source_splits

Existing split names combined to form the pool.

Type:: list[str]

group_by

Columns used for group-based splitting.

Type:: str | list[str] | None

stratify_by

Columns used for stratified splitting.

Type:: str | list[str] | None

Configure cross-validation for a single FeatureSet.

Parameters:

fs (str | FeatureSet) – FeatureSet (or its node ID/label) to apply cross-validation to.
source_splits (list[str]) – Existing splits of fs to pool before folding. For example, source_splits=[‘train’, ‘val’] merges both splits prior to sampling.
group_by (str | list[str] | None, optional) – Optional tag keys used for group-based splitting. Mutually exclusive with stratify_by. Defaults to None.
stratify_by (str | list[str] | None, optional) – Optional tag keys used for stratified splitting. Mutually exclusive with group_by. Defaults to None.
train_split_name (str, optional) – Split label that should receive each fold’s training partition. Defaults to train.
val_split_name (str, optional) – Split label that should receive each fold’s validation partition. Defaults to val.
val_size (float | None, optional) – Explicit validation proportion for each fold. If None, computed as 1 / n_folds. Defaults to None.

Raises:

ValueError – If configuration references missing splits or invalid sizes.

Callbacks#

class EarlyStopping(*, monitor: str = 'val_loss', mode: Literal['min', 'max'] = 'min', patience: int = 5, min_delta: float = 0.0, restore_best: bool = False, reducer: Literal['mean', 'sum', 'last', 'first'] = 'last', label: str | None = None, execution_order: int = 1)[source]#

Bases: Callback

Stop training when a monitored metric stops improving.

This callback monitors a named metric from the MetricStore and stops training if no improvement is observed for a given number of epochs (patience). Improvement is determined by the mode parameter: “min” expects the metric to decrease, “max” expects it to increase.

When triggered, EarlyStopping calls phase.request_stop() which sets a flag that the training loop checks at the end of each epoch.

When restore_best=True, the model state is restored to the epoch with the best monitored metric value at the end of the phase. If the phase has a Checkpointing callback configured, its saved states are used for restoration. Otherwise, EarlyStopping manages its own in-memory state snapshots as a fallback.

Example

>>> # Stop if validation loss doesn't improve for 5 epochs
>>> phase.add_callback(  
...     EarlyStopping(monitor="val_loss", patience=5)
... )

>>> # Stop and restore model to best epoch
>>> phase.add_callback(  
...     EarlyStopping(
...         monitor="val_loss",
...         patience=5,
...         restore_best=True,
...     )
... )

__init__(*, monitor: str = 'val_loss', mode: Literal['min', 'max'] = 'min', patience: int = 5, min_delta: float = 0.0, restore_best: bool = False, reducer: Literal['mean', 'sum', 'last', 'first'] = 'last', label: str | None = None, execution_order: int = 1) → None[source]#

Initialize an EarlyStopping callback.

Parameters:

monitor (str, optional) – Name of the metric to monitor. Must match a metric name that is logged into the MetricStore during training (e.g. by a EvalLossMetric or custom MetricCallback). Defaults to “val_loss”.
mode (Literal["min", "max"], optional) – Whether the monitored metric should be minimized or maximized. Defaults to “min”.
patience (int, optional) – Number of epochs with no improvement after which training will be stopped. Defaults to 5.
min_delta (float, optional) – Minimum change in the monitored metric to qualify as an improvement. Defaults to 0.0.
restore_best (bool, optional) – Whether to restore the model state to the epoch with the best monitored metric value at the end of the phase. If the phase has a Checkpointing callback, its saved states are used. Otherwise, in-memory snapshots are managed automatically. Defaults to False.
reducer (str, optional) – If the monitor metric is produced more than one per epoch, reducer defines how to aggregate all values in that epoch. Typically, monitor is produced at most once per epoch and this argument is not used.
label (str | None, optional) – Stable identifier for this callback. Defaults to “EarlyStopping”.
execution_order (int, optional) – Used for execution ordering or multiple callbacks, where higher values are executed later than lower values. This value should be greater than the callback that produces the monitor metric. Unless you manually modified the other callback execution orders, a value of 1 is fine.

property best_epoch: int | None#: The epoch index that achieved the best metric value, or None.

property best_value: float | None#: The best observed metric value so far.

classmethod from_config(config: dict) → EarlyStopping[source]#: Construct from config data.

get_config() → dict[str, Any][source]#: Return configuration details required to reconstruct this callback.

property monitor: str#: The name of the monitored metric.

on_epoch_end(*, experiment: Experiment, phase: ExperimentPhase, exec_ctx: ExecutionContext, results: PhaseResults | None = None) → CallbackResult | None[source]#: Check monitored metric and request stop if patience is exceeded.

on_phase_end(*, experiment: Experiment, phase: ExperimentPhase, results: PhaseResults | None = None) → EarlyStoppingResult | None[source]#: Restore best model state (if enabled) and return summary result.

on_phase_start(*, experiment: Experiment, phase: ExperimentPhase, results: PhaseResults | None = None) → None[source]#: Reset internal state at the start of each phase.

property restore_best: bool#: Whether model state restoration to the best epoch is enabled.

property stopped_epoch: int | None#: The epoch at which training was stopped, or None if not triggered.

class EvalLossMetric(*, loss: AppliedLoss, reducer: Literal['sum', 'mean'] = 'mean', name: str = 'val_loss')[source]#

Bases: EvaluationMetric

Extracts a scalar loss from an Evaluation callback and logs it as a metric.

This built-in EvaluationMetric reads the results of its parent Evaluation callback and extracts the aggregated loss value for a specific node. The extracted value is logged as a named metric (default: “val_loss”) into the MetricStore.

Example

Below showcases created a tracked “val_loss” during a training phase.

>>> mse_loss = AppliedLoss(...)  
>>> val_metric = EvalLossMetric(
...     loss=mse_loss, name="val_loss"
... )  
>>> eval_cb = Evaluation(  
...     eval_phase=eval_phase, metrics=[val_metric]
... )
>>> phase.add_callback(eval_cb)  

__init__(*, loss: AppliedLoss, reducer: Literal['sum', 'mean'] = 'mean', name: str = 'val_loss') → None[source]#

Initialize a ValidationLossMetric.

Parameters:

loss (AppliedLoss) – An applied loss to track. Will be appended to the parent EvalPhase in Evaluation, if not already present.
reducer (Literal["sum", "mean"], optional) – How to aggregate per-batch losses from the evaluation into a single scalar. Defaults to “mean”.
name (str, optional) – The metric name to log under. Defaults to “val_loss”.
mode (Literal["min", "max"], optional) – Whether lower or higher values are better. Defaults to “min”.

extract(*, eval_result: EvaluationCallbackResult, exec_ctx: ExecutionContext | None = None) → MetricResult | None[source]#: Extract aggregated loss from the Evaluation result.

classmethod from_config(config: dict) → EvalLossMetric[source]#: Construct from config data.

get_config() → dict[str, Any][source]#: Return configuration details required to reconstruct this callback.

Modeling#

class ModelGraph(nodes: list[str | GraphNode] | None, optimizer: Optimizer | None = None, label: str = 'model-graph', *, ctx: ExperimentContext | None = None, register: bool = True)[source]#

Bases: Configurable, Stateful

Directed acyclic graph orchestrating GraphNode nodes.

label

User-defined identifier for the graph.

Type:: str

nodes

Registered nodes keyed by GraphNode.node_id.

Type:: dict[str, GraphNode]

_optimizer

Optional graph-level optimizer.

Type:: Optimizer | None

__init__(nodes: list[str | GraphNode] | None, optimizer: Optimizer | None = None, label: str = 'model-graph', *, ctx: ExperimentContext | None = None, register: bool = True)[source]#

Initialize a ModelGraph from a list of modular nodes and an optional global optimizer.

Parameters:

nodes (list[str | GraphNode], optional) – A list of GraphNodes (e.g., ModeNode) or their labels to construct a ModelGraph around. If None, all registered GraphNodes in the active ExperimentContext are used.
optimizer (Optional[Optimizer], optional) – A shared optimizer to use for all nodes that require one. If provided, the graph will ensure that all such stages have a matching backend and override any stage-level optimizers. If not provided, each stage that requires an optimizer must define one locally.
label (str, optional) – Optional label to assign to this ModelGraph instance. Defaults to “model-graph-0”.
ctx (ExperimentContext, optional) – The ExperimentContext this ModelGraph should exist within. If None, uses the active ExperimentContext. Defaults to None.
register (bool, optional) – Used only for de-serialization.

Raises:

ValueError – If duplicate node labels are provided or if graph connectivity is invalid.
RuntimeError – If required optimizers are missing or if backends are incompatible.

add_node(node: GraphNode) → ModelGraph[source]#

Add a new node to the graph.

Description:: This modifies graph structure only; no existing node states are reset or copied. The added node must already be registered in the ExperimentContext.

Parameters:: node (GraphNode) – Node to add.
Returns:: self
Return type:: ModelGraph
Raises:: ValueError – If a node with the same node_id already exists.

property backend: Backend | None#

The shared backend of this ModelGraph.

Description:: A ModelGraph’s backend is only defined if a global optimizer if used. If the graph consists of mixed-backend nodes, None is returned.

Returns:: The backend of the global optimizer, if defined. Otherwise, returns None.
Return type:: Backend | None

build(*, force: bool = False)[source]#

Build the ModelGraph by initializing all underlying models and optimizers.

Parameters:: force (bool, optional) – If the graph is already built it will not be rebuilt unless force=True. Defaults to False.

eval_step(ctx: ExecutionContext, losses: list[AppliedLoss], *, active_nodes: list[str | GraphNode] | None = None)[source]#

Execute a single evaluation step for the ModelGraph.

Description:: Performs a forward-only pass through the graph, computes all applicable losses, and records outputs and losses into the ExecutionContext. No gradients are tracked and no optimizers are stepped.

Parameters:

ctx (ExecutionContext) – Execution context containing inputs, outputs, and loss storage.
losses (list[AppliedLoss]) – Losses to compute during evaluation.
active_nodes (list[str | GraphNode] | None, optional) – Optional subset of nodes to execute. All required upstream dependencies are included automatically.

fit_step(ctx: ExecutionContext, losses: list[AppliedLoss] | None = None, *, active_nodes: list[str | GraphNode] | None = None, freeze_after_fit: bool = True)[source]#

Fit batch-fit nodes in topological order.

Description:

Iterates through active nodes in topological order. Nodes that implement the Fittable protocol and are not frozen will have fit_step() called. All other forwardable nodes perform a forward-only pass to propagate outputs downstream.

After fitting, nodes are optionally frozen to prevent interference during subsequent gradient-based training phases.

Parameters:

ctx (ExecutionContext) – Execution context containing full-dataset inputs.
losses (list[AppliedLoss] | None, optional) – Optional losses to compute after fitting (for metrics only).
active_nodes (list[str | GraphNode] | None, optional) – Optional subset of nodes to fit. If None, all nodes in the graph are considered.
freeze_after_fit (bool, optional) – Whether to freeze fitted nodes after completion. Defaults to True.

forward(inputs: dict[tuple[str, FeatureSetReference], TForward], *, active_nodes: list[str | GraphNode] | None = None) → dict[str, Batch][source]#

Execute a forward pass through the ModelGraph.

Parameters:

inputs (dict[tuple[str, FeatureSetReference], TForward]) – Mapping of (head_node_id, upstream_featureset_ref) -> TForward. Keys must correspond to head nodes in this graph (nodes whose upstream refs are FeatureSetReferences). A head node may have multiple inputs if it consumes multiple FeatureSets.
active_nodes (list[str | GraphNode] | None, optional) – Optional subset of nodes to execute. If provided, only these nodes (and any required upstream dependencies within this graph) are executed. If None, all nodes in the graph are executed.

Returns:

Mapping of node_id -> output for every executed node (typically all nodes, but may be restricted by active_nodes). Tail-node outputs can be obtained by filtering this dict with self.tail_nodes.

Return type:

dict[str, TForward]

freeze(nodes: list[str, GraphNode] | None = None)[source]#

Sets the trainable state of nodes to frozen.

Parameters:: nodes (list[str, GraphNode] | None) – A list of node IDs, labels, or instances. All specified nodes will have their internal state set to frozen, preventing training mutation. If None, all nodes in this graph will be frozen.

classmethod from_config(config: dict[str, Any], *, register: bool = True) → ModelGraph[source]#

Reconstructs a ModelGraph from configuration details.

This does not restore state information of any underlying models or optimizers.

property frozen_nodes: dict[str, GraphNode]#: All trainable but frozen nodes in this graph, keyed by node_id.

get_config() → dict[str, Any][source]#

Retrieve the configuration details of this ModelGraph instance.

This does not contain state information of any underlying models or optimizers.

get_optimizer_parameters() → dict[str, Any][source]#

State of current global optimizer (if defined).

Returns a dict with the set of node_ids actually contributing parameters, and backend specific fields:

“backend”: Backend,
“contributing_nodes”: set[str],
“parameters”: list[torch.nn.Parameter], # PyTorch only
“variables”: list[tf.Variable], # TensorFlow only

get_state() → dict[str, Any][source]#

Return serialized state for all nodes and global optimizer.

Returns:

Snapshot capturing node state, optimizer: state, and build metadata.

Return type:

dict[str, Any]

property head_nodes: dict[str, GraphNode]#

Head nodes of the ModelGraph.

Description:: Head nodes are GraphNodes whose inputs originate directly from FeatureSets (i.e., they have no upstream GraphNode dependencies).

Returns:: Mapping of node_id to GraphNode for all head nodes.
Return type:: dict[str, GraphNode]

insert_node_after(new_node: GraphNode, *, upstream: str | GraphNode) → ModelGraph[source]#

Insert a node after an existing node.

Description:: Inserts a new node after an existing GraphNode. A new output connection is added between the existing node and the new node. All other output connection are left undisturbed. The new node will only receive input from the existing node.

Parameters:

new_node (GraphNode) – New node instance to be inserted.
upstream (str | GraphNode) – Node ID, label, or instance of an existing ModelGraph node. The new_node will be inserted downstream of this node.

Returns:

self

Return type:

insert_node_before(new_node: GraphNode, *, downstream: str | GraphNode) → ModelGraph[source]#

Insert a node before an existing node.

Description:: Inserts a new node before an existing GraphNode. All inputs to the existing node are attached to the new node. The existing node will now only receive input from the new node.

Parameters:

new_node (GraphNode) – New node instance to be inserted.
downstream (str | GraphNode) – Node ID, label, or instance of an existing ModelGraph node. The new_node will be inserted upstream of this node.

Returns:

self

Return type:

insert_node_between(new_node: GraphNode, *, upstream: str | GraphNode, downstream: str | GraphNode) → ModelGraph[source]#

Insert a node between two, already connected, nodes.

Description:: Insert a new node between a connection of two existing node. The old connection (upstream -> downstream) is replaced with (upstream -> new_node -> downstream). An error will be thrown if the existing nodes are not already connected.

Parameters:

new_node (GraphNode) – New node instance to be inserted.
upstream (str | GraphNode) – Node ID, label, or instance of an existing ModelGraph node. The new_node will be inserted downstream of this node.
downstream (str | GraphNode) – Node ID, label, or instance of an existing ModelGraph node. The new_node will be inserted upstream of this node.

Returns:

self

Return type:

property is_built: bool#

Return whether the graph has been initialized.

Returns:: True once build() has run successfully.
Return type:: bool

classmethod load(filepath: Path, *, allow_packaged_code: bool = False, overwrite: bool = False) → ModelGraph[source]#

Load a FeaturModelGrapheSet from file.

Parameters:

filepath (Path) – File location of a previously saved ModelGraph.
allow_packaged_code – bool Whether bundled code execution is allowed.
overwrite (bool) – Whether to replace any colliding node registrations in ExperimentContext If False, new IDs are assigned to the reloaded nodes comprising the graph. Otherwise, any collision are overwritten with the saved nodes. Defaults to False. It is recommended to only reload a ModelGraph into a new/empty ExperimentContext.

Returns:

The reloaded ModelGraph.

Return type:

property node_labels: set[str]#: Returns the set of unique node labels in this ModelGraph.

property nodes: dict[str, GraphNode]#: All GraphNode in the ModelGraph, keyed by node_id.

remove_node(node: str | GraphNode) → ModelGraph[source]#

Remove an existing node from the graph.

Description:: The existing node is removed an all connections are updated. Any nodes downstream of node will re-route inputs to all nodes that previously provided input to node.

Parameters:: node (GraphNode) – Node ID, label, or instance of an existing ModelGraph node to be removed.

Example

Removing an existing single-input, single-output node.

Given: A -> B -> C

Remove: B

Result: A -> C
Removing an existing multi-input, single-output node.

Given: [A, B] -> C -> D

Remove: C

Result: [A, B] -> D

Note that D must be able to accept multiple inputs or an error will be thrown.
Removing an existing single-input, multi-output node.

Given: A -> B -> [C, D]

Remove: B

Result: A -> [C, D]

Note that A must be able to accept multiple outputs or an error will be thrown.

Returns:: self
Return type:: ModelGraph

replace_node(old_node: str | GraphNode, new_node: GraphNode) → ModelGraph[source]#

Replace an existing node while preserving all upstream and downstream connections.

Description:: Connectivity of the graph is preserved, and learned state of the other nodes is unaffected. The replaced node’s state is replaced with the new node.

Parameters:

old_node (str | GraphNode) – Existing node in the graph to be replaced. Provided argument value can be the existing node’s label, ID, or the actual node instance. The state of the existing node is not changed; the graph connections are simply redirected to the new_node.
new_node (GraphNode) – New node instance to take the spot of old_node.

Returns:

self

Return type:

restore_checkpoint(filepath: Path) → ModelGraph[source]#

Restore ModelGraph state from checkpoint.

Parameters:: filepath (Path) – File location of a previously saved ModelGraph checkpoint.
Returns:: The ModelGraph restored to the checkpoint state.
Return type:: self

save(filepath: Path, *, overwrite: bool = False) → Path[source]#

Serializes this ModelGraph to the specified filepath.

Parameters:

filepath (Path) – File location to save to. Note that the suffix may be overwritten to enforce the ModularML file extension schema.
overwrite (bool, optional) – Whether to overwrite any existing file at the save location. Defaults to False.

Returns:

The actual filepath to write the ModelGraph is saved.

Return type:

Path

save_checkpoint(filepath: Path, *, overwrite: bool = False, meta: dict[str, Any] | None = None) → Path[source]#

Save full ModelGraph state.

Parameters:

filepath (Path) – File location to save to. Note that the suffix may be overwritten to enforce the ModularML file extension schema.
overwrite (bool, optional) – Whether to overwrite any existing file at the save location. Defaults to False.
meta (dict[str, Any], optional) – Additional meta data to attach to the checkpoint. Must be pickle-able.

Returns:

Final path of saved ModelGraph checkpoint.

Return type:

Path

set_state(state: dict[str, Any]) → None[source]#

Restore node and optimizer state from get_state().

Parameters:: state (dict[str, Any]) – Snapshot previously generated by get_state().

property tail_nodes: dict[str, GraphNode]#

Tail nodes of the graph.

Description:: Tail nodes are GraphNodes whose outputs are not consumed by any other GraphNode in the ModelGraph.

Returns:: Mapping of node_id to GraphNode for all tail nodes.
Return type:: dict[str, GraphNode]

train_step(ctx: ExecutionContext, losses: list[AppliedLoss], *, active_nodes: list[str | GraphNode] | None = None)[source]#

Execute a single training step for the ModelGraph.

Behavior depends on whether a global optimizer is attached:

If self._optimizer is None:
Stage-wise training is performed. Each ModelNode executes its own train_step() in topological order using its local optimizer.
If self._optimizer is not None:
Graph-wise training is performed. A single forward pass is executed across the graph, all losses are computed, a single backward pass is performed, and the global optimizer is stepped once.

Parameters:

ctx (ExecutionContext) – Execution context containing inputs, outputs, and loss storage.
losses (list[AppliedLoss]) – All losses defined for this phase. Losses are filtered internally to the nodes they apply to.
active_nodes (list[str | GraphNode] | None, optional) –
Optional subset of nodes to be excuted in the forward pass. If provided, all upstream dependencies are included automatically. Otherwise, all nodes in the graph are executed.

Note that this does not set the trainable state of the nodes, only on which nodes a forward pass is called. Use freeze() and unfreeze to set the trainable state of graph nodes.

unfreeze(nodes: list[str, GraphNode] | None = None)[source]#

Sets the trainable state of nodes to unfrozen.

Parameters:: nodes (list[str, GraphNode] | None) – A list of node IDs, labels, or instances. All specified nodes will have their internal state set to unfrozen, allowing training. If None, all nodes in this graph will be unfrozen.

visualize(*, show_features: bool = False, show_targets: bool = False, show_tags: bool = False, show_frozen: bool = True, show_splits: bool = False)[source]#

Displays a mermaid diagram for this FeatureSet.

Parameters:

show_features (bool, optional) – Show feature columns on head nodes. Defaults to False.
show_targets (bool, optional) – Show target columns on head nodes. Defaults to False.
show_tags (bool | str, optional) – Show tags columns on head nodes. Defaults to False.
show_frozen (bool, optional) – Show frozen state (label text and dimmed styling) on ModelNodes Defaults to True.
show_splits (bool, optional) – Show available splits on FeatureSet nodes. Defaults to False.

class ModelNode(label: str, model: BaseModel | Any, upstream_ref: ExperimentNode | ExperimentNodeReference, optimizer: Optimizer | None = None, *, node_id: str | None = None, register: bool = True)[source]#

Bases: ComputeNode

Single learnable or static stage inside a ModelGraph.

_model

Wrapped backend model implementation.

Type:: BaseModel

_optimizer

Optional optimizer coordinating gradient steps.

Type:: Optimizer | None

_freeze

Flag indicating whether training is disabled.

Type:: bool

__init__(label: str, model: BaseModel | Any, upstream_ref: ExperimentNode | ExperimentNodeReference, optimizer: Optimizer | None = None, *, node_id: str | None = None, register: bool = True)[source]#

Initialize a ModelNode.

Parameters:

label (str) – Unique name identifying this stage within the model graph.
model (Union[BaseModel, Any]) – A backend-specific model instance or config.
upstream_ref (ExperimentReference) – Reference to the upstream node.
optimizer (Optional[Optimizer]) – Optimizer to use during training (optional).
node_id (str, optional) – Used only for de-serialization.
register (bool, optional) – Used only for de-serialization.

property backend: Backend#

Returns the backend associated with the wrapped model.

Returns:: TORCH, TENSORFLOW, SCIKIT, …
Return type:: Backend

build_model(input_shape: tuple[int, ...] | None = None, output_shape: tuple[int, ...] | None = None, *, force: bool = False)[source]#

Build the ModelNode by initializing the internal BaseModel and optimizer.

Parameters:

input_shape (tuple[int, ...] | None, optional) – Input shape to construct this model with. Defaults to None.
output_shape (tuple[int, ...] | None, optional) – Output shape to construct this model with. If not provided, the BaseModel must be capable of inferring it internally or during construction. Defaults to None.
force (bool, optional) – If model is already instantiated it will not be re-instantiated unless force=True. Defaults to False.

Notes

For PyTorch and TensorFlow, optimizers are built after the model is initialized.
Scikit-learn models typically do not require external optimizers.
This method assumes that shape inference and merge logic (if needed) has already been resolved upstream by the ModelGraph.

eval_step(ctx: ExecutionContext, losses: list[AppliedLoss] | None = None)[source]#

Performs an evaluation step (forward pass and loss computation) for this stage.

Only callable if this stage is frozen. No gradient tracking is performed.

Parameters:

ctx (ExecutionContext) – Context (input/output data) for the given execution step.
losses (list[AppliedLoss]) – Optional list of losses to be applied in this execution step.

Raises:

RuntimeError – If stage is not frozen.

fit_step(ctx: ExecutionContext, losses: list[AppliedLoss] | None = None)[source]#

Fits this node on complete data (for batch-fit scikit-learn models).

Calls the underlying model’s .fit(X, y) method using the full dataset provided in the execution context. After fitting, a forward pass is performed to record outputs for downstream nodes.

Parameters:

ctx (ExecutionContext) – Context containing full-dataset inputs.
losses (list[AppliedLoss] | None) – Optional losses to compute after fitting (for metrics only).

Raises:

RuntimeError – If this node is frozen.

forward_single(batch: Batch, **kwargs) → Batch[source]#

forward_single(roles: RoleData, **kwargs) → RoleData

forward_single(data: SampleData, **kwargs) → SampleData

Performs a forward pass through the model using SampleData.

This method preserves raw tensor outputs to maintain backend autograd support. It returns a SampleData object keyed by output roles containing model predictions.

Parameters:

x (SampleData | RoleData | Batch) – Input data to the model.
**kwargs – Any additional keyword arguments to provide to BaseModel.forward

Returns:

Outputs from the model. Output type matches input.

Return type:

SampleData | RoleData | Batch

freeze()[source]#: Freezes this node (prevents training updates).

classmethod from_config(config: dict[str, Any], *, register: bool = True) → ModelNode[source]#

Reconstructs a ModelNode from configuration details.

This does not restore state information of the underlying model or optimizer.

get_config() → dict[str, Any][source]#

Retrieve the configuration details of this ModelNode instance.

This does not contain state information of the underlying model or optimizer.

get_state() → dict[str, Any][source]#

Return serialized state for the node, model, and optimizer.

Returns:: Snapshot captured for set_state().
Return type:: dict[str, Any]

property input_shape: tuple[int, ...]#

Return the model’s input tensor shape.

Returns:: Expected feature tensor shape.
Return type:: tuple[int, …]

property is_built: bool#

Checks if the model has been built (i.e., instantiated with input/output shape).

Returns:: True if built, False otherwise.
Return type:: bool

property is_frozen: bool#

Indicates whether this stage is frozen (not trainable).

Returns:: True if frozen, False if trainable.
Return type:: bool

property max_upstream_refs: int#

Return the maximum number of allowed upstream references.

Returns:: Always 1 because ModelNode has a single input.
Return type:: int

property model: BaseModel#

Return the wrapped backend model instance.

Returns:: Backend-specific implementation.
Return type:: BaseModel

property output_shape: tuple[int, ...]#

Return the model’s output tensor shape.

Returns:: Output tensor shape.
Return type:: tuple[int, …]

set_state(state: dict[str, Any]) → None[source]#

Restore runtime state from get_state() output.

Parameters:: state (dict[str, Any]) – Serialized node data.

train_step(ctx: ExecutionContext, losses: list[AppliedLoss])[source]#

Performs a training step (forward, loss, backward, optimizer step) for this stage.

Only callable if this stage has an optimizer and is not frozen. Otherwise, training must be delegated to ModelGraph.

Parameters:

ctx (ExecutionContext) – Context (input/output data) for the given execution step.
losses (list[AppliedLoss]) – List of losses to be applied in this execution step.

Raises:

RuntimeError – If stage is frozen or optimizer is missing.

unfreeze()[source]#: Unfreezes this node (allows training updates).

class ConcatNode(label: str, upstream_refs: list[ExperimentNode | ExperimentNodeReference], concat_axis: int = 0, *, concat_axis_targets: StrategyType | ExperimentNodeReference = -1, concat_axis_tags: StrategyType | ExperimentNodeReference = -1, pad_inputs: bool = False, pad_mode: str | PadMode = 'constant', pad_value: float = 0.0, node_id: str | None = None, register: bool = True)[source]#

Bases: MergeNode

A merge stage that concatenates multiple inputs along a specified axis.

Description:

This stage merges tensors by concatenating them along a specified axis. It supports automatic padding of non-concat dimensions to align shapes, allowing for flexible merging even when inputs vary in size. Padding behavior can be controlled via mode (e.g., ‘constant’, ‘reflect’, ‘replicate’) and value.

For the targets and tags domains, non-concatenation merge strategies can be used instead of axis-based concatenation. For example, “first” selects targets from the first input, “mean” computes an element-wise average, or an ExperimentNodeReference can be passed to select targets from a specific upstream input.

label

Unique identifier for this node.

Type:: str

upstream_refs

Upstream node references from which inputs will be received.

Type:: list[ExperimentNode | ExperimentNodeReference]

concat_axis

The axis along which to concatenate feature inputs.

Type:: int

target_strategy

Strategy for merging targets. An int means concatenation along that axis; a MergeStrategy applies an aggregation; an ExperimentNodeReference selects targets from a specific upstream input.

Type:: int | MergeStrategy | ExperimentNodeReference

tags_strategy

Strategy for merging tags (same semantics as target_strategy).

Type:: int | MergeStrategy | ExperimentNodeReference

pad_inputs

Whether to pad inputs before merging. Defaults to False.

Type:: bool, optional

pad_mode

Padding mode (‘constant’, ‘reflect’, ‘replicate’, etc.). Defaults to ‘constant’.

Type:: PadMode, optional

pad_value

Value to use for constant padding. Defaults to 0.0.

Type:: float, optional

__init__(label: str, upstream_refs: list[ExperimentNode | ExperimentNodeReference], concat_axis: int = 0, *, concat_axis_targets: StrategyType | ExperimentNodeReference = -1, concat_axis_tags: StrategyType | ExperimentNodeReference = -1, pad_inputs: bool = False, pad_mode: str | PadMode = 'constant', pad_value: float = 0.0, node_id: str | None = None, register: bool = True)[source]#

Initialize a ConcatNode.

Parameters:

label (str) – Unique identifier for this node.
upstream_refs (list[ExperimentNode | ExperimentNodeReference]) – Upstream node references from which inputs will be received.
concat_axis (int) –
The axis along which to concatenate feature inputs. Does not include the batch dimension. That is, for shape (with batch) of (32,1,16), axis=0 refers to “1”. The examples below omit the batch dimension.
- axis=0: concat along features. Example: (1, 16) + (1, 16) -> (2, 16).
- axis > 1: concat within features. Example: (1, 16, 16) + (1, 16, 8) -> (1, 16, 24). All data must have at least concat_axis dimensions.
- axis=-1: concat along the last axis of all data. Example: (1, 16, 16) + (1, 16, 8) -> (1, 16, 24).
concat_axis_targets (int | str | MergeStrategy | ExperimentNodeReference) –
Strategy for merging the targets domain. Accepts:
- int: concatenate along this axis (same semantics as concat_axis). Defaults to -1 (last axis).
- str or MergeStrategy: apply an aggregation strategy. Supported values are "first", "last", and "mean".
- ExperimentNodeReference: select targets from the upstream input matching this reference.
concat_axis_tags (int | str | MergeStrategy | ExperimentNodeReference) – Strategy for merging the tags domain. Same semantics as concat_axis_targets. Defaults to -1 (last axis).
pad_inputs (bool, optional) – Whether to pad inputs before merging. Defaults to False.
pad_mode (PadMode, optional) – Padding mode; one of {“constant”, “reflect”, “replicate”, “circular”}. Defaults to “constant”.
pad_value (float, optional) – Value to use for constant padding. Defaults to 0.0.
node_id (str, optional) – Used only for de-serialization.
register (bool, optional) – Used only for de-serialization.

apply_merge(values: list[Any], *, includes_batch_dim: bool = True, fmt: DataFormat | None = None, domain: str = 'features') → Any[source]#

Concatenate input tensors along the configured axis.

Description:: Optionally pads the inputs to align non-concat dimensions before applying backend-specific concatenation. This method handles axis-based concatenation only. Non-concat strategies (e.g., “first”, “mean”) are handled by _merge_sample_data() before reaching this method.

Parameters:

values (list[Any]) – A list of backend-specific tensors to be merged.
includes_batch_dim (bool) – Whether the input values have a batch dimension. Defaults to True.
fmt (DataFormat | None) – The data format expected for the returned tensor. If None, the data format will be inferred from the backend property. Defaults to None.
domain (str, optional) – The domain in which the data belongs. This allows for domain-specific merge logic (e.g., different concat axes for each domain).

Returns:

Concatenated tensor, in the specified format.

Return type:

Any

classmethod from_config(config: dict, *, register: bool = True) → MergeNode[source]#

Construct a ConcatNode from a configuration dictionary.

Parameters:

config (dict[str, Any]) – Configuration details. Keys must be strings.
register (bool) – Whether to register the reconstructed node.

Returns:

Reconstructed ConcatNode.

Return type:

Callback

get_config() → dict[str, Any][source]#

Return configuration details required to reconstruct this node.

Returns:: Configuration used to reconstruct the node. Keys must be strings.
Return type:: dict[str, Any]

property tags_axis: int#

Tags concatenation axis (only valid when tags_strategy is int).

Raises:: TypeError – If the tags strategy is not int-based.

property target_axis: int#

Target concatenation axis (only valid when target_strategy is int).

Raises:: TypeError – If the target strategy is not int-based.

Bases: object

Backend-agnostic wrapper around loss functions used in model training.

Description:: Supports built-in loss classes from PyTorch and TensorFlow along with custom callables. Provides serialization, backend normalization, and config/state handling to slot into ModularML training workflows.

name

Lowercase name of the selected loss when applicable.

Type:: str | None

cls

Backing loss class when using backend libraries.

Type:: type | None

fn

Direct loss callable if provided.

Type:: Callable | None

reduction

Reduction passed to backend constructors.

Type:: str

kwargs

Keyword arguments used for construction.

Type:: dict[str, Any] | None

_backend

Resolved backend enum.

Type:: Backend | None

Initialize the loss wrapper with a name, class, callable, or factory.

Parameters:

loss (str | type | Callable | None) – Loss identifier, class, or callable; mutually exclusive with factory.
loss_kwargs (dict[str, Any] | None) – Keyword arguments passed to loss construction.
backend (Backend | None) – Backend enum; required when loss is a string or factory.
reduction (str) – Reduction argument forwarded to backend constructors.
factory (Callable | None) – Callable producing a loss instance when invoked.

Returns:

This initializer does not return a value.

Return type:

None

Raises:

ValueError – If both loss and factory are provided, backend inference fails, or inputs are invalid.
LossError – If initialization does not provide a valid loss definition.

property allowed_keywords: list[str]#

List valid keyword arguments for the currently built loss callable.

Returns:: Argument names accepted by _callable.
Return type:: list[str]
Raises:: RuntimeError – If the loss has not been built yet.

property backend: Backend#

Backend configured for this Loss.

Returns:: Resolved backend enum.
Return type:: Backend
Raises:: LossError – If the backend has not been set.

build(*, backend: Backend | None = None, force_rebuild: bool = False, **kwargs)[source]#

Instantiate the loss callable if not already provided.

Parameters:

backend (Backend | None) – Backend to use, overriding the inferred backend if provided.
force_rebuild (bool) – Whether to rebuild even if already instantiated.
**kwargs (Any) – Additional keyword arguments forwarded to loss construction.

Raises:

LossError – If rebuilding an already built loss without force_rebuild.
ValueError – If backend mismatches occur or no backend is set before building.
BackendNotSupportedError – If the backend lacks constructor support.

classmethod from_config(config: dict[str, Any]) → Loss[source]#

Construct a Loss from configuration.

Parameters:: config (dict[str, Any]) – Serialized configuration dictionary.
Returns:: Rebuilt loss wrapper instance.
Return type:: Loss

get_config() → dict[str, Any][source]#

Return configuration required to reconstruct this loss wrapper.

Returns:: Serialized configuration capturing loss definition.
Return type:: dict[str, Any]

property is_built: bool#

Whether the underlying loss callable has been instantiated.

Returns:: True if _callable is available.
Return type:: bool

classmethod load(filepath: Path, *, allow_packaged_code: bool = False) → Loss[source]#

Load a loss wrapper from disk.

Parameters:

filepath (Path) – Path to a serialized loss artifact.
allow_packaged_code (bool) – Whether packaged code execution is allowed.

Returns:

Reloaded loss instance.

Return type:

Loss

save(filepath: Path, *, overwrite: bool = False) → Path[source]#

Serialize this loss wrapper to disk.

Parameters:

filepath (Path) – Destination path; suffix may be adjusted to match ModularML conventions.
overwrite (bool) – Whether to overwrite existing files.

Returns:

Actual file path written by the serializer.

Return type:

Path

class AppliedLoss(loss: Loss, on: str | ModelNode, inputs: list[ReferenceLike] | dict[str, ReferenceLike], *, weight: float = 1.0, label: str | None = None)[source]#

Bases: Summarizable

Bind a Loss to a ModelNode with resolved inputs.

loss

Wrapped Loss instance evaluated for each execution step.

Type:: Loss

weight

Scalar multiplier applied to computed loss values before aggregation.

Type:: float

label

Friendly label used when logging summaries for this applied loss.

Type:: str | None

node_id

Identifier of the ModelNode targeted by this loss.

Type:: str

inputs

Mapping of loss argument names to runtime references.

Type:: dict[str, ReferenceLike]

__init__(loss: Loss, on: str | ModelNode, inputs: list[ReferenceLike] | dict[str, ReferenceLike], *, weight: float = 1.0, label: str | None = None)[source]#

Define a Loss applied on a specified ModelNode.

Parameters:

loss (Loss) – Loss instance to apply for each execution step.
on (str | ModelNode) – Node label/ID or ModelNode object indicating where the loss attaches.
inputs (list[ReferenceLike] | dict[str, ReferenceLike]) – Positional or keyword references resolved as loss arguments.
weight (float) – Scalar multiplier applied to the computed loss value.
label (str | None) – Optional label used when logging or summarizing this applied loss.

Raises:

TypeError – If on is not a ModelNode or string, or inputs has an unsupported type.

property backend: Backend#

Backend used by the underlying Loss.

Returns:: Backend declared on the wrapped Loss.
Return type:: Backend

compute(ctx: ExecutionContext) → Any[source]#

Compute the weighted loss for a single execution step.

Parameters:

ctx (ExecutionContext) – Execution context supplying model outputs and upstream batches.

Returns:

Backend-specific scalar/tensor representing the weighted loss value.

Return type:

Any

Raises:

BackendMismatchError – If the Loss backend differs from the ModelNode backend.
ResolutionError – If any configured reference cannot be resolved from the execution context.

classmethod from_config(config: dict[str, Any]) → AppliedLoss[source]#

Construct an AppliedLoss from configuration.

Parameters:: config (dict[str, Any]) – Dictionary produced by get_config().
Returns:: Rehydrated applied loss instance.
Return type:: AppliedLoss

get_config() → dict[str, Any][source]#

Return configuration required to reconstruct this applied loss.

Returns:: Serialized dictionary capturing loss, node, inputs, weight, and label.
Return type:: dict[str, Any]

Bases: Configurable, Stateful

Backend-agnostic optimizer wrapper that lazily constructs backend objects.

Description:: Supports initialization from optimizer names, classes, or callables while normalizing backend selection for PyTorch and TensorFlow.

name

Lowercase name of the optimizer when known.

Type:: str | None

cls

Backend optimizer class resolved via introspection.

Type:: type | None

kwargs

Keyword arguments applied when instantiating the optimizer.

Type:: dict[str, Any]

_backend

Resolved backend enum for this optimizer.

Type:: Backend | None

instance

Concrete optimizer object once built.

Type:: Any | None

parameters

Stored model parameters used for PyTorch optimizer construction.

Type:: Any | None

_pending_state

Deferred backend state restored after Optimizer.build().

Type:: dict[str, Any] | None

Initialize the optimizer wrapper from a name, class, or factory.

Parameters:

opt (str | type | None) – Optimizer name or class, mutually exclusive with factory.
opt_kwargs (dict[str, Any] | None) – Keyword arguments provided when instantiating the optimizer.
factory (Callable | None) – Callable returning an optimizer when invoked during Optimizer.build().
backend (Backend | None) – Backend enum required for name/factory initialization.

Raises:

ValueError – If arguments conflict, backend is missing when required, or inputs are invalid.
TypeError – If opt is neither a string nor a type.

property backend: Backend | None#

Backend enum associated with this optimizer.

Returns:: Resolved backend value if available.
Return type:: Backend | None

build(*, parameters: Any | None = None, backend: Backend | None = None, force_rebuild: bool = False)[source]#

Instantiate the backend optimizer if not already provided.

Parameters:

parameters (Any | None) – Trainable parameters required when building PyTorch optimizers.
backend (Backend | None) – Backend override enforced before construction.
force_rebuild (bool) – Whether to rebuild even if the optimizer is already instantiated.

Raises:

OptimizerNotSetError – If attempting to rebuild without force_rebuild.
ValueError – If backend validations fail or parameters are missing for PyTorch.
BackendNotSupportedError – If an unsupported backend is requested.
RuntimeError – If the initialization mode is unsupported.

classmethod from_config(config: dict[str, Any]) → Optimizer[source]#

Instantiate an optimizer from serialized configuration.

Parameters:: config (dict[str, Any]) – Configuration dictionary produced by get_config().
Returns:: Reconstructed optimizer wrapper.
Return type:: Optimizer

classmethod from_factory(factory: Callable, *, backend: Backend) → Optimizer[source]#

Instantiate an Optimizer directly from a factory and backend.

Parameters:

factory (Callable) – Callable that produces an optimizer instance.
backend (Backend) – Backend enum applied to the optimizer.

Returns:

New wrapper configured to call the provided factory.

Return type:

Optimizer

get_config() → dict[str, Any][source]#

Return the serialized configuration for this optimizer.

Returns:: Dictionary capturing optimizer definition for reconstruction.
Return type:: dict[str, Any]

get_state() → dict[str, Any][source]#

Capture serialized optimizer state including backend internals.

Returns:: State payload containing build flag and backend data.
Return type:: dict[str, Any]

property is_built: bool#: Whether the backend optimizer instance has been constructed.

classmethod load(filepath: Path, *, allow_packaged_code: bool = False) → Optimizer[source]#

Load a serialized optimizer from disk.

Parameters:

filepath (Path) – Path to a serialized optimizer artifact.
allow_packaged_code (bool) – Whether bundled code execution is permitted.

Returns:

Reloaded optimizer instance.

Return type:

Optimizer

save(filepath: Path, *, overwrite: bool = False) → Path[source]#

Serialize the optimizer to disk using the built-in serializer.

Parameters:

filepath (Path) – Destination path; suffix may be adjusted to match conventions.
overwrite (bool) – Whether to overwrite an existing artifact.

Returns:

Actual file path written by the serializer.

Return type:

Path

set_state(state: dict[str, Any]) → None[source]#

Store serialized state for later restoration during build().

Parameters:: state (dict[str, Any]) – Serialized optimizer state produced by get_state().

step(grads=None, variables=None)[source]#

Perform a backend-specific optimizer step.

Parameters:

grads (Any | None) – Gradient tensors required by TensorFlow optimizers.
variables (Any | None) – Trainable variables paired with grads in TensorFlow.

Raises:

OptimizerNotSetError – If the optimizer has not been built.
ValueError – If TensorFlow requires gradients or variables that are missing.
BackendNotSupportedError – If the backend is unsupported.

zero_grad()[source]#

Reset accumulated gradients on the backend optimizer.

Raises:

OptimizerNotSetError – If the optimizer has not been built.
BackendNotSupportedError – If the backend is unsupported.

class BaseModel(backend: str | Backend, **init_args: Any)[source]#

Bases: Configurable, Stateful, ABC

Abstract base class for backend-agnostic ModularML models.

_backend

Normalized backend enum powering the model.

Type:: Backend

_built

Flag indicating whether build() completed.

Type:: bool

_init_args

Keyword arguments stored for config round-trips via get_config().

Type:: dict[str, Any]

__init__(backend: str | Backend, **init_args: Any)[source]#

Initialize the model with a backend and keyword arguments.

Parameters:

backend (str | Backend) – Backend identifier or enum value used to normalize _backend.
**init_args (Any) – Keyword arguments cached for config serialization so subclasses can reconstruct themselves.

property backend: Backend#: Return the Backend powering this model.

abstract build(input_shape: tuple[int, ...] | None = None, output_shape: tuple[int, ...] | None = None, *, force: bool = False)[source]#

Build backend layers for the model implementation.

Parameters:

input_shape (tuple[int, ...] | None) – Per-sample input shape, excluding batch dimension, if known.
output_shape (tuple[int, ...] | None) – Per-sample output shape, excluding batch dimension, if known.
force (bool) – Whether to rebuild even if the model is already constructed.

abstract forward(*args, **kwargs)[source]#: Run a forward pass for the underlying backend.

classmethod from_config(config: dict[str, Any]) → BaseModel[source]#

Construct a BaseModel from serialized configuration.

Description:: Instantiates a new model with the provided config. Learned weights are not restored; call set_state() afterwards.

Parameters:: config (dict[str, Any]) – Configuration emitted by get_config().
Returns:: Newly constructed model instance.
Return type:: BaseModel

get_config() → dict[str, Any][source]#

Return configuration required to re-instantiate this model.

Returns:

Serializable configuration capturing the model: class, backend, and build metadata.

Return type:

dict[str, Any]

get_init_args() → dict[str, Any][source]#

Return keyword arguments needed to reconstruct this model.

Description:: Subclasses may override this to control config serialization, but caching init_args automatically works for most models.

Returns:: Copy of initialization keyword arguments.
Return type:: dict[str, Any]

get_state() → dict[str, Any][source]#: Return the learned state for serialization.

abstract get_weights() → dict[str, Any][source]#: Return backend weights as pure-Python or numpy data.

abstract property input_shape: tuple[int, ...] | None#

Return the expected per-sample input shape for the model.

Returns:

Shape tuple without the batch dimension,: or None when the shape is not yet known.

Return type:

tuple[int, …] | None

property is_built: bool#: Return True when the model has been built and shapes resolved.

classmethod load(filepath: Path, *, allow_packaged_code: bool = False) → BaseModel[source]#

Load a serialized BaseModel from disk.

Parameters:

filepath (Path) – Path pointing to a saved model.
allow_packaged_code (bool) – Whether packaged code execution is permitted for user-defined artifacts.

Returns:

Reloaded model instance.

Return type:

BaseModel

abstract property output_shape: tuple[int, ...] | None#

Return the per-sample output shape produced by the model.

Returns:

Shape tuple without the batch dimension,: or None when not yet known.

Return type:

tuple[int, …] | None

save(filepath: Path, *, overwrite: bool = False) → Path[source]#

Serialize this model to disk using the serializer.

Parameters:

filepath (Path) – Destination path. The suffix may be adjusted to follow ModularML naming conventions.
overwrite (bool) – Whether to overwrite an existing artifact.

Returns:

Actual file path written by the serializer.

Return type:

Path

set_state(state: dict[str, Any]) → None[source]#

Restore learned state produced by get_state().

Parameters:: state (dict[str, Any]) – State dictionary containing weights.

abstract set_weights(weights: dict[str, Any]) → None[source]#

Restore the weights previously produced by get_weights().

Parameters:: weights (dict[str, Any]) – Dictionary returned by get_weights().

class TorchBaseModel(*args: Any, **kwargs: Any)[source]#

Bases: BaseModel, Module, ABC

Base class for ModularML-native PyTorch models.

Description:: Intended for framework-owned PyTorch architectures such as modularml.models.torch.SequentialMLP. User-defined modules can subclass this base, although TorchModelWrapper is usually simpler for existing torch.nn.Module graphs.

__init__(**init_args: Any)[source]#: Initialize the PyTorch + BaseModel inheritance chain.

get_weights() → dict[str, ndarray][source]#: Return PyTorch tensors as numpy arrays via state_dict().

set_weights(weights: dict[str, ndarray]) → None[source]#: Restore numpy-based weights produced by get_weights().

class TensorflowBaseModel(**init_args: Any)[source]#

Bases: BaseModel, ABC

Base class for ModularML-native TensorFlow/Keras models.

Description:: Designed for framework-owned Keras implementations. User-defined tf.keras.Model objects can subclass this base, though TensorflowModelWrapper is typically simpler.

__init__(**init_args: Any)[source]#: Initialize TensorFlow dependencies and call BaseModel.

get_weights() → dict[str, np.ndarray][source]#: Return model weights as numpy arrays keyed by variable name.

set_weights(weights: dict[str, np.ndarray]) → None[source]#: Restore numpy-based weights produced by get_weights().

FeatureSets#

class FeatureSet(label: str, collection: SampleCollection, **kwargs)[source]#

Bases: ExperimentNode, SplitMixin, SampleCollectionMixin

Unified FeatureSet backed by a single SampleCollection.

Each representation (e.g., “raw”, “transformed”) lives within the same SampleCollection rather than as separate sub-collections. Splits store indices into this collection and may specify which representation(s) to use when retrieving data.

__init__(label: str, collection: SampleCollection, **kwargs)[source]#

Initialize a ExperimentNode with a name and register.

Parameters:

label (str) – Unique identifier for this node.
node_id (str, optional) – Used only for de-serialization.
register (bool, optional) – Used only for de-serialization.

add_split(split: FeatureSetView)[source]#

Adds a new FeatureSetView.

Parameters:: split (FeatureSetView) – The new view to add.

property available_splits: list[str]#

All available splits.

Returns:: Available split names.
Return type:: list[str]

clear_splits() → None[source]#: Removes all previously defined splits.

Create a reference to a single column in this FeatureSet.

Parameters:

column (str | None) – Fully-qualified column name to reference (e.g. “features.voltage.raw”).
feature (str | None) – Feature-domain selector. Domain prefix may be omitted.
target (str | None) – Target-domain selector. Domain prefix may be omitted.
tag (str | None) – Tag-domain selector. Domain prefix may be omitted.
rep (str | None) – Default representation suffix applied when a selector omits a representation.

Returns:

FeatureSetColumnReference

copy(*, label: str | None = None, share_raw_data_buffer: bool = True, restore_splits: bool = False, restore_scalers: bool = False, register: bool = False) → FeatureSet[source]#

Create a copy of this FeatureSet with optional state restoration.

Description:

Constructs a new FeatureSet instance based on this one. By default, the PyArrow buffers of the raw data columns are shared (zero-copy). Splitters and scalers can optionally be re-applied to the new instance, producing non-shared transform data columns.

This method does not mutate the original FeatureSet.

Parameters:

label (str | None, optional) – Label for the new FeatureSet. If None, appends “_copy” to the current label.
share_raw_data_buffer (bool, optional) – If True, PyArrow buffers of raw data columns are shared between the original and copied FeatureSet (zero-copy). If False, a deep copy of the Arrow table is created.
restore_splits (bool, optional) – If True, all stored SplitterRecords are re-applied to the new FeatureSet to regenerate splits.
restore_scalers (bool, optional) – If True, all ScalerRecords are re-applied in order to regenerate transformed representations. If True, restore_splits must be also be enabled.
register (bool, optional) – Whether to register the copied FeatureSet in the ExperimentContext registry. If True, a new node ID will be generated for the copied instance. Must be True if restore_scalers or restore_splits are enabled.

Returns:

A new FeatureSet instance.

Return type:

Raises:

ValueError – If restore_scalers=True but restore_splits=False and scalers depend on split-specific fitting.

Fit a scaler to selected columns and apply the transform to the entire FeatureSet.

Description:

Flattens the selected columns (2D), optionally merging axes (merged_axes).
Fits the provided scaler to samples from fit_to_split (or all samples).
Transforms all samples.
Unflattens results back to original shapes.
Stores a new representation (REP_TRANSFORMED) for each key.
Appends a ScalerRecord to the transform log.

Parameters:

scaler (Scaler | Any) – A Scaler instance or any sklearn-like object implementing fit() and transform(). If not a Scaler, it is wrapped as Scaler(scaler=<obj>).
domain (str) – One of {“features”, “targets”, “tags”}.
keys (str | list[str] | None, optional) – Column names within the domain to transform. If None, all domain keys are used. Defaults to None.
fit_to_split (str | None, optional) – Split name from which to fit the scaler (e.g., “train”). If None, the scaler is fit to all samples. Defaults to None.
merged_axes (int | tuple[int]) – Axes whose sizes are merged into a single dimension. If None, no axes are merged. An Error will be thrown if the resulting shape is not 2-dimensional.

Raises:

ValueError –
- If flattening with merged_axes does not result in 2D data. - If representation or domain keys do not exist.
TypeError – If scaler is not a valid Scaler or sklearn-transformer-like object.

classmethod from_config(config: dict[str, Any], *, register: bool = True) → FeatureSet[source]#: Instantiate an empty FeatureSet.

Construct a FeatureSet from a pandas DataFrame (column-wise storage).

Description:

Converts a DataFrame into a column-oriented Arrow table that matches the ModularML SampleSchema convention. Each domain (features, targets, tags) becomes a Struct column in the final Arrow table.

If group_by are provided, all rows sharing the same group key are aggregated into a single sample row, with feature and target columns stored as array-valued sequences (e.g., np.ndarray or list).

Parameters:

label (str) – Name to assign to this FeatureSet.
df (pd.DataFrame) – Input pandas DataFrame containing raw experimental or measurement data.
feature_cols (str | list[str]) – Column name(s) in df to be used as feature variables.
target_cols (str | list[str]) – Column name(s) in df to be used as target variables.
group_by (str | list[str] | None, optional) – One or more column names defining group boundaries. Each group becomes one sample (row) in the final table, with grouped columns stored as lists. If None, each original DataFrame row is treated as a sample.
tag_cols (str | list[str] | None, optional) – Column name(s) corresponding to identifying or categorical metadata (e.g., cell ID, protocol, SOC). Defaults to None.
sort_by (str | list[str] | None, optional) – Column name(s) used to sort rows within each group before aggregation. Only used when group_by is specified; ignored otherwise. Defaults to None.

Returns:

A new Arrow-backed FeatureSet whose table follows the ModularML SampleSchema convention.

Return type:

Example

FeatureSet construction from a Pandas dataframe:

>>> fs = FeatureSet.from_pandas( 
...     label="PulseData",
...     df=raw_df,
...     feature_cols=["voltage", "current"],
...     target_cols=["soh"],
...     group_by=["cell_id", "cycle_index"],
...     tag_cols=["temperature", "cell_id"],
...     sort_by="timestamp",
... )

classmethod from_dict(label: str, data: dict[str, Sequence[Any]], feature_keys: str | list[str], target_keys: str | list[str], tag_keys: str | list[str] | None = None) → FeatureSet[source]#

Construct a FeatureSet from a Python dictionary of column data.

Description:

Converts a dictionary of lists/arrays into a column-oriented Arrow table following the ModularML SampleSchema convention. Each key in the input dictionary corresponds to a column name, and values are list-like sequences of equal length representing sample data.

Unlike from_pandas, this constructor assumes that each list element already represents a complete sample (i.e., no grouping is applied).

Parameters:

label (str) – Label to assign to this FeatureSet.
data (dict[str, Sequence[Any]]) – A mapping from column names to list-like column data. Each list must have the same length, corresponding to the total number of samples.
feature_keys (str | list[str]) – Column name(s) in data to be used as features.
target_keys (str | list[str]) – Column name(s) in data to be used as targets.
tag_keys (str | list[str] | None, optional) – Column name(s) corresponding to identifying or categorical metadata (e.g., cell ID, protocol, SOC). Defaults to None.

Returns:

A new Arrow-backed FeatureSet containing all provided columns, organized into standardized domains (features, targets, and tags).

Return type:

Raises:

ValueError – If required keys are missing, or column lengths are inconsistent.

Example

FeatureSet construction via a dict of data:

>>> fs = FeatureSet.from_dict( 
...     label="CycleData",
...     data={
...         "voltage": [[3.1, 3.2, 3.3], [3.2, 3.3, 3.4]],
...         "current": [[1.0, 1.1, 1.0], [0.9, 1.0, 1.1]],
...         "soh": [0.95, 0.93],
...         "cell_id": ["A1", "A2"],
...     },
...     feature_keys=["voltage", "current"],
...     target_keys=["soh"],
...     tag_keys=["cell_id"],
...     )

Construct a FeatureSet from a pandas DataFrame (column-wise storage).

Description:

Converts a DataFrame into a column-oriented Arrow table that matches the ModularML SampleSchema convention. Each domain (features, targets, tags) becomes a Struct column in the final Arrow table.

If group_by are provided, all rows sharing the same group key are aggregated into a single sample row, with feature and target columns stored as array-valued sequences (e.g., np.ndarray or list).

Parameters:

label (str) – Name to assign to this FeatureSet.
df (pd.DataFrame) – Input pandas DataFrame containing raw experimental or measurement data.
feature_cols (str | list[str]) – Column name(s) in df to be used as feature variables.
target_cols (str | list[str]) – Column name(s) in df to be used as target variables.
group_by (str | list[str] | None, optional) – One or more column names defining group boundaries. Each group becomes one sample (row) in the final table, with grouped columns stored as lists. If None, each original DataFrame row is treated as a sample.
tag_cols (str | list[str] | None, optional) – Column name(s) corresponding to identifying or categorical metadata (e.g., cell ID, protocol, SOC). Defaults to None.
sort_by (str | list[str] | None, optional) – Column name(s) used to sort rows within each group before aggregation. Only used when group_by is specified; ignored otherwise. Defaults to None.

Returns:

A new Arrow-backed FeatureSet whose table follows the ModularML SampleSchema convention.

Return type:

Example

FeatureSet construction from a Pandas dataframe:

>>> fs = FeatureSet.from_pandas( 
...     label="PulseData",
...     df=raw_df,
...     feature_cols=["voltage", "current"],
...     target_cols=["soh"],
...     group_by=["cell_id", "cycle_index"],
...     tag_cols=["temperature", "cell_id"],
...     sort_by="timestamp",
... )

classmethod from_pyarrow_table(label: str, table: pa.Table, schema: SampleSchema | None = None, **kwargs) → FeatureSet[source]#

Construct a new FeatureSet from an existing PyArrow table.

Parameters:

label (str) – Label to assign to this FeatureSet.
table (pa.Table) – Table to build FeatureSet around.
schema (SampleSchema | None) – PyArrow schema to use for SampleCollection.
kwargs – Additional keyword arguments.

Returns:

New FeatureSet instance

Return type:

get_config() → dict[str, Any][source]#: Structural configuration of the FeatureSet and parents.

get_split(split_name: str) → FeatureSetView[source]#

Gets the specified split, rebuilding the FeatureSetView if necessary.

Description:: If new columns (e.g. transformed representations) have been added to the FeatureSet after the split was defined, the cached FeatureSetView is re-created to include all current columns while preserving row indices.

Parameters:: split_name (str) – Name of the split to retrieve.
Returns:: A no-copy, row-wise view of the FeatureSet.
Return type:: FeatureSetView

get_state() → dict[str, Any][source]#: Runtime state (i.e., PyArrow table and records) of the FeatureSet.

iter_scalers_on_cols(*, domain: str, columns: str | list[str]) → Iterable[Scaler][source]#

Iterate over scaler objects applicable to the provided columns.

Description:: Resolves the scaler dependency chain for inverse scaling and returns scaler objects in the correct order for calling inverse_transform.

Parameters:

domain (str) – One of {“features”, “targets”, “tags”}.
columns (str | list[str]) – Column keys within the domain (no domain prefix, no rep suffix).

Returns:

Scalers ordered newest to oldest.

Return type:

list[Scaler]

Raises:

ValueError – If scalers cannot be safely resolved.

classmethod load(filepath: Path, *, allow_packaged_code: bool = False, overwrite: bool = False) → FeatureSet[source]#

Load a FeatureSet from file.

Parameters:

filepath (Path) – File location of a previously saved FeatureSet.
allow_packaged_code – bool Whether bundled code execution is allowed.
overwrite (bool) – Whether to replace any colliding node registrations in ExperimentContext If False, a new node_id is assigned to the reloaded FeatureSet. Otherwise, the existing FeatureSet is removed from the ExperimentContext registry. Defaults to False.

Returns:

The reloaded FeatureSet.

Return type: