Reference

Core

The model class that turns a pytorch_symbolic graph into a stateful, forecast-capable reservoir computer, plus the symbolic-input helpers used to declare its entry points.

core

Core Model

This module provides the main model class and input helpers for building ESN architectures using the pytorch_symbolic library.

CLASS	DESCRIPTION
ESNModel	Extended SymbolicModel with ESN-specific methods for forecasting and reservoir state management.
ReservoirFeatureExtractor	nn.Sequential-friendly adapter that packages reservoir layers as a drop-in feature extractor for ordinary torch.nn pipelines.
Input	Alias for pytorch_symbolic.Input for defining model inputs.

FUNCTION	DESCRIPTION
reservoir_input	Convenience constructor for a per-feature symbolic input tensor with a placeholder time dimension. Preferred over hand-crafting ps.Input((T, F)) since the time dimension is purely a tracing hint inside reservoir models.

Examples:

Building a simple ESN:

>>> from resdag.core import ESNModel, reservoir_input
>>> from resdag.layers import ESNLayer
>>> from resdag.layers.readouts import CGReadoutLayer
>>>
>>> inp = reservoir_input(3)             # equivalent to ps.Input((1, 3))
>>> reservoir = ESNLayer(200, feedback_size=3)(inp)
>>> readout = CGReadoutLayer(200, 3)(reservoir)
>>> model = ESNModel(inp, readout)

Multi-input model:

>>> feedback = reservoir_input(3)
>>> driver = reservoir_input(5)
>>> reservoir = ESNLayer(200, feedback_size=3, input_size=5)(feedback, driver)
>>> readout = CGReadoutLayer(200, 3)(reservoir)
>>> model = ESNModel([feedback, driver], readout)

See Also

resdag.models : Premade ESN architectures. resdag.training.ESNTrainer : Trainer for fitting readouts.

Input module-attribute

Input = ps.Input

ESNModel

ESNModel(*args: Any, **kwargs: Any)

Bases: SymbolicModel

Echo State Network model with forecasting and state management.

This class extends pytorch_symbolic.SymbolicModel with ESN-specific functionality including:

• Time series forecasting with warmup and autoregressive generation
• Reservoir state management (reset, get, set)
• Model persistence (save/load)
• Architecture visualization

The model inherits all standard torch.nn.Module functionality and the summary() method from pytorch_symbolic.

Pass symbolic inputs and outputs to the constructor (see pytorch_symbolic.SymbolicModel).

ATTRIBUTE	DESCRIPTION
inputs	List of model input tensors. TYPE: list
outputs	List of model output tensors. TYPE: list
output_shape	Shape(s) of model outputs. TYPE: torch.Size or tuple of torch.Size

Examples:

Create and use a simple ESN:

>>> import pytorch_symbolic as ps
>>> from resdag.core import ESNModel
>>> from resdag.layers import ESNLayer
>>> from resdag.layers.readouts import CGReadoutLayer
>>>
>>> inp = ps.Input((100, 3))
>>> reservoir = ESNLayer(200, feedback_size=3)(inp)
>>> readout = CGReadoutLayer(200, 3)(reservoir)
>>> model = ESNModel(inp, readout)
>>>
>>> # Forward pass
>>> x = torch.randn(4, 100, 3)
>>> y = model(x)
>>> print(y.shape)
torch.Size([4, 100, 3])

Forecasting with warmup:

>>> warmup_data = torch.randn(1, 50, 3)
>>> predictions = model.forecast(warmup_data, horizon=100)
>>> print(predictions.shape)
torch.Size([1, 100, 3])

See Also

pytorch_symbolic.SymbolicModel : Parent class. BaseReservoirLayer : Reservoir layer component. ESNTrainer : Trainer for fitting readout layers.

pytorch_symbolic.SymbolicModel propagates shapes by running each layer's real forward on a torch.rand placeholder (batch size 1). For a stateful :class:~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer that leaves the layer holding a saturated trace state at batch=1. A fresh model used directly at batch=1 (e.g. y = model(x) or standard forecasting) would silently continue from this junk instead of a clean zero state — batch sizes != 1 auto-resize to zeros and escape it, but batch=1 direct calls do not. Resetting here, once, right after the graph is built, makes every freshly constructed model start from a clean state regardless of how it is first called.

Source code in src/resdag/core/model.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    """Build the symbolic model and clear reservoir trace residue.

    ``pytorch_symbolic.SymbolicModel`` propagates shapes by running each
    layer's real ``forward`` on a ``torch.rand`` placeholder (batch size 1).
    For a stateful :class:`~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer`
    that leaves the layer holding a *saturated* trace state at ``batch=1``.
    A fresh model used directly at ``batch=1`` (e.g. ``y = model(x)`` or
    standard forecasting) would silently continue from this junk instead of a
    clean zero state — batch sizes ``!= 1`` auto-resize to zeros and escape
    it, but ``batch=1`` direct calls do not. Resetting here, once, right after
    the graph is built, makes every freshly constructed model start from a
    clean state regardless of how it is first called.
    """
    super().__init__(*args, **kwargs)
    self.reset_reservoirs()

reset_reservoirs

reset_reservoirs() -> None

Reset all reservoir layer states to zero.

This clears the internal hidden states of all :class:BaseReservoirLayer modules in the model, preparing it for a new sequence.

Examples:

>>> model.reset_reservoirs()
>>> output = model(new_sequence)

Source code in src/resdag/core/model.py

def reset_reservoirs(self) -> None:
    """
    Reset all reservoir layer states to zero.

    This clears the internal hidden states of all :class:`BaseReservoirLayer`
    modules in the model, preparing it for a new sequence.

    Examples
    --------
    >>> model.reset_reservoirs()
    >>> output = model(new_sequence)
    """
    for module in self.modules():
        if isinstance(module, BaseReservoirLayer):
            module.reset_state()

set_random_reservoir_states

set_random_reservoir_states(batch_size: int | None = None, device: device | None = None, dtype: dtype | None = None) -> None

Set random states of all reservoir layers.

PARAMETER	DESCRIPTION
batch_size	If provided, lazily initialise each reservoir's state with this batch size before filling it with random values. TYPE: int DEFAULT: None
device	Target device for lazy initialisation. TYPE: device DEFAULT: None
dtype	Target dtype for lazy initialisation. TYPE: dtype DEFAULT: None

Examples:

>>> model.set_random_reservoir_states()                # state must exist
>>> model.set_random_reservoir_states(batch_size=4)    # lazy

Source code in src/resdag/core/model.py

def set_random_reservoir_states(
    self,
    batch_size: int | None = None,
    device: torch.device | None = None,
    dtype: torch.dtype | None = None,
) -> None:
    """
    Set random states of all reservoir layers.

    Parameters
    ----------
    batch_size : int, optional
        If provided, lazily initialise each reservoir's state with this
        batch size before filling it with random values.
    device : torch.device, optional
        Target device for lazy initialisation.
    dtype : torch.dtype, optional
        Target dtype for lazy initialisation.

    Examples
    --------
    >>> model.set_random_reservoir_states()                # state must exist
    >>> model.set_random_reservoir_states(batch_size=4)    # lazy
    """
    for module in self.modules():
        if isinstance(module, BaseReservoirLayer):
            module.set_random_state(batch_size=batch_size, device=device, dtype=dtype)

get_reservoir_states

get_reservoir_states() -> dict[str, Tensor]

Get current states of all reservoir layers.

RETURNS	DESCRIPTION
dict of str to torch.Tensor	Dictionary mapping layer names to their state tensors. Only includes reservoirs with non-None states.

Examples:

>>> states = model.get_reservoir_states()
>>> for name, state in states.items():
...     print(f"{name}: {state.shape}")

Source code in src/resdag/core/model.py

def get_reservoir_states(self) -> dict[str, torch.Tensor]:
    """
    Get current states of all reservoir layers.

    Returns
    -------
    dict of str to torch.Tensor
        Dictionary mapping layer names to their state tensors.
        Only includes reservoirs with non-None states.

    Examples
    --------
    >>> states = model.get_reservoir_states()
    >>> for name, state in states.items():
    ...     print(f"{name}: {state.shape}")
    """
    states = {}
    for name, module in self.named_modules():
        if isinstance(module, BaseReservoirLayer) and module.state is not None:
            states[name] = module.state.clone()
    return states

set_reservoir_states

set_reservoir_states(states: dict[str, Tensor], strict: bool = True) -> None

Set states of reservoir layers.

PARAMETER	DESCRIPTION
states	Dictionary mapping layer names to state tensors. Names should match those returned by :meth:get_reservoir_states. TYPE: dict of str to torch.Tensor
strict	If True, raise a KeyError when the provided states dict is missing entries for any reservoir in the model or contains keys that don't match any reservoir. Set to False to silently ignore both kinds of mismatch (legacy behaviour). TYPE: bool DEFAULT: True

Notes

Each restored tensor is routed through :meth:~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.set_state, which clones the tensor and validates its shape against the target cell's contract (e.g. the 2-D (batch, reservoir_size) layout of an ESN). Before validation the tensor is coerced to the target reservoir's reference device and dtype (its first floating-point parameter/buffer) so that a state saved on one device/dtype is not silently re-zeroed by _maybe_init_state on the next forward pass — the canonical save-on-GPU / load-on-CPU round-trip therefore preserves the warmed-up state values.

RAISES	DESCRIPTION
KeyError	If strict=True and the keys of states do not exactly match the set of reservoir layer names in the model.
ValueError	If a restored tensor does not match the target cell's state-shape contract. The error names the cell class and the offending shape.

WARNS	DESCRIPTION
UserWarning	If a restored state had to be moved to a different device or cast to a different dtype to match its target reservoir.

Examples:

>>> states = model.get_reservoir_states()
>>> # ... do something ...
>>> model.set_reservoir_states(states)  # Restore states

Source code in src/resdag/core/model.py

def set_reservoir_states(
    self,
    states: dict[str, torch.Tensor],
    strict: bool = True,
) -> None:
    """
    Set states of reservoir layers.

    Parameters
    ----------
    states : dict of str to torch.Tensor
        Dictionary mapping layer names to state tensors.
        Names should match those returned by :meth:`get_reservoir_states`.
    strict : bool, default=True
        If ``True``, raise a ``KeyError`` when the provided ``states``
        dict is missing entries for any reservoir in the model or
        contains keys that don't match any reservoir.  Set to ``False``
        to silently ignore both kinds of mismatch (legacy behaviour).

    Notes
    -----
    Each restored tensor is routed through
    :meth:`~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.set_state`,
    which clones the tensor and validates its shape against the target
    cell's contract (e.g. the 2-D ``(batch, reservoir_size)`` layout of an
    ESN).  Before validation the tensor is coerced to the target
    reservoir's reference device and dtype (its first floating-point
    parameter/buffer) so that a state saved on one device/dtype is not
    silently re-zeroed by ``_maybe_init_state`` on the next forward pass —
    the canonical save-on-GPU / load-on-CPU round-trip therefore preserves
    the warmed-up state values.

    Raises
    ------
    KeyError
        If ``strict=True`` and the keys of ``states`` do not exactly
        match the set of reservoir layer names in the model.
    ValueError
        If a restored tensor does not match the target cell's state-shape
        contract.  The error names the cell class and the offending shape.

    Warns
    -----
    UserWarning
        If a restored state had to be moved to a different device or cast
        to a different dtype to match its target reservoir.

    Examples
    --------
    >>> states = model.get_reservoir_states()
    >>> # ... do something ...
    >>> model.set_reservoir_states(states)  # Restore states
    """
    import warnings

    reservoir_names = {
        name for name, module in self.named_modules() if isinstance(module, BaseReservoirLayer)
    }
    provided = set(states.keys())

    if strict:
        missing = reservoir_names - provided
        extra = provided - reservoir_names
        if missing or extra:
            parts: list[str] = []
            if missing:
                parts.append(f"missing keys for reservoirs: {sorted(missing)}")
            if extra:
                parts.append(f"unexpected keys not matching any reservoir: {sorted(extra)}")
            raise KeyError("set_reservoir_states(strict=True): " + "; ".join(parts))

    for name, module in self.named_modules():
        if isinstance(module, BaseReservoirLayer) and name in states:
            state = states[name]
            ref_device, ref_dtype = module.reference_device_dtype()
            if state.device != ref_device or state.dtype != ref_dtype:
                warnings.warn(
                    f"set_reservoir_states: state for reservoir '{name}' was "
                    f"coerced from (device={state.device}, dtype={state.dtype}) to "
                    f"(device={ref_device}, dtype={ref_dtype}) to match the "
                    f"target reservoir.",
                    UserWarning,
                    stacklevel=2,
                )
                state = state.to(device=ref_device, dtype=ref_dtype)
            # set_state clones and validates the tensor via the cell's
            # state-shape contract.
            module.set_state(state)

save

save(path: str | Path, include_states: bool = False, **metadata: Any) -> None

Save model weights and optionally reservoir states.

PARAMETER	DESCRIPTION
path	File path to save the model. Parent directories are created if they don't exist. TYPE: str or Path
include_states	If True, also save current reservoir states. Only reservoirs whose state is non-None are persisted (see :meth:get_reservoir_states); reservoirs that are still lazily-uninitialised at save time — e.g. a model saved before any warmup — are omitted from the checkpoint. :meth:load restores the saved states tolerantly, so this partial-key checkpoint round-trips without error. TYPE: bool DEFAULT: False
**metadata	Additional metadata to store with the model (e.g., training info). TYPE: Any DEFAULT: {}

Examples:

Save model weights only:

>>> model.save("model.pt")

Save with states and metadata:

>>> model.save(
...     "checkpoint.pt",
...     include_states=True,
...     epoch=10,
...     loss=0.05
... )

See Also

load : Load model from file.

Source code in src/resdag/core/model.py

def save(
    self,
    path: str | Path,
    include_states: bool = False,
    **metadata: Any,
) -> None:
    """
    Save model weights and optionally reservoir states.

    Parameters
    ----------
    path : str or Path
        File path to save the model. Parent directories are created
        if they don't exist.
    include_states : bool, default=False
        If True, also save current reservoir states. Only reservoirs whose
        state is non-``None`` are persisted (see
        :meth:`get_reservoir_states`); reservoirs that are still
        lazily-uninitialised at save time — e.g. a model saved before any
        warmup — are omitted from the checkpoint. :meth:`load` restores the
        saved states tolerantly, so this partial-key checkpoint round-trips
        without error.
    **metadata
        Additional metadata to store with the model (e.g., training info).

    Examples
    --------
    Save model weights only:

    >>> model.save("model.pt")

    Save with states and metadata:

    >>> model.save(
    ...     "checkpoint.pt",
    ...     include_states=True,
    ...     epoch=10,
    ...     loss=0.05
    ... )

    See Also
    --------
    load : Load model from file.
    """
    path = Path(path)

    # Create parent directories if they don't exist
    path.parent.mkdir(parents=True, exist_ok=True)

    save_dict = {
        "state_dict": self.state_dict(),
        "metadata": metadata,
    }

    if include_states:
        save_dict["reservoir_states"] = self.get_reservoir_states()

    torch.save(save_dict, path)

load

load(path: str | Path, strict: bool = True, load_states: bool = False) -> None

Load model weights from file.

PARAMETER	DESCRIPTION
path	File path to load from. TYPE: str or Path
strict	If True, strictly enforce that state_dict keys match. TYPE: bool DEFAULT: True
load_states	If True, also load reservoir states if available. States are restored tolerantly (strict=False): because :meth:save only persists reservoirs whose state is non-None, a checkpoint may legitimately omit some reservoir keys (e.g. one saved before warmup). Reservoirs absent from the checkpoint keep their current state and no KeyError is raised. TYPE: bool DEFAULT: False

WARNS	DESCRIPTION
UserWarning	If load_states=True but no states are found in checkpoint.

Examples:

>>> model.load("model.pt")
>>> model.load("checkpoint.pt", load_states=True)

See Also

save : Save model to file. load_from_file : Class method for loading.

Source code in src/resdag/core/model.py

def load(
    self,
    path: str | Path,
    strict: bool = True,
    load_states: bool = False,
) -> None:
    """
    Load model weights from file.

    Parameters
    ----------
    path : str or Path
        File path to load from.
    strict : bool, default=True
        If True, strictly enforce that state_dict keys match.
    load_states : bool, default=False
        If True, also load reservoir states if available. States are
        restored tolerantly (``strict=False``): because
        :meth:`save` only persists reservoirs whose state is non-``None``,
        a checkpoint may legitimately omit some reservoir keys (e.g. one
        saved before warmup). Reservoirs absent from the checkpoint keep
        their current state and no ``KeyError`` is raised.

    Warns
    -----
    UserWarning
        If ``load_states=True`` but no states are found in checkpoint.

    Examples
    --------
    >>> model.load("model.pt")
    >>> model.load("checkpoint.pt", load_states=True)

    See Also
    --------
    save : Save model to file.
    load_from_file : Class method for loading.
    """
    import warnings

    path = Path(path)
    checkpoint = torch.load(path, weights_only=False)

    if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint:
        raise ValueError(
            f"{path} is not a save() state-dict checkpoint. If it was written "
            f"by save_full() (or torch.save(model)), use ESNModel.load_full()."
        )

    self.load_state_dict(checkpoint["state_dict"], strict=strict)

    if load_states:
        if "reservoir_states" in checkpoint:
            # save(include_states=True) only persists reservoirs whose state
            # is non-None (see save / get_reservoir_states), so a checkpoint
            # written before warmup legitimately omits some reservoir keys.
            # Restore tolerantly (strict=False) to mirror that asymmetry —
            # reservoirs missing from the checkpoint keep their current
            # (typically None / lazily-initialised) state.
            self.set_reservoir_states(checkpoint["reservoir_states"], strict=False)
        else:
            warnings.warn(
                "load_states=True but no reservoir states found in checkpoint", UserWarning
            )

load_from_file classmethod

load_from_file(path: str | Path, model: 'ESNModel | None' = None, strict: bool = True, load_states: bool = False) -> 'ESNModel'

Load weights into an existing model instance.

This is a convenience class method that loads state dict into a pre-constructed model.

PARAMETER	DESCRIPTION
path	File path to load from. TYPE: str or Path
model	Model instance to load weights into. Required. TYPE: ESNModel DEFAULT: None
strict	If True, strictly enforce state_dict key matching. TYPE: bool DEFAULT: True
load_states	If True, also load reservoir states. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
ESNModel	The model instance with loaded weights.

RAISES	DESCRIPTION
ValueError	If model is None.

Examples:

>>> model = create_my_model()  # Create architecture
>>> model = ESNModel.load_from_file("weights.pt", model=model)

Source code in src/resdag/core/model.py

@classmethod
def load_from_file(
    cls,
    path: str | Path,
    model: "ESNModel | None" = None,
    strict: bool = True,
    load_states: bool = False,
) -> "ESNModel":
    """
    Load weights into an existing model instance.

    This is a convenience class method that loads state dict into
    a pre-constructed model.

    Parameters
    ----------
    path : str or Path
        File path to load from.
    model : ESNModel
        Model instance to load weights into. Required.
    strict : bool, default=True
        If True, strictly enforce state_dict key matching.
    load_states : bool, default=False
        If True, also load reservoir states.

    Returns
    -------
    ESNModel
        The model instance with loaded weights.

    Raises
    ------
    ValueError
        If ``model`` is None.

    Examples
    --------
    >>> model = create_my_model()  # Create architecture
    >>> model = ESNModel.load_from_file("weights.pt", model=model)
    """
    if model is None:
        raise ValueError("model argument is required")

    model.load(path, strict=strict, load_states=load_states)
    return model

save_full

save_full(path: str | Path, **metadata: Any) -> None

Serialize the entire model — architecture, weights, and reservoir states — to a single file.

Unlike :meth:save, which stores only the state_dict (so the architecture must be re-created before :meth:load), this pickles the whole model object, the pytorch_symbolic graph included. Restore it with :meth:load_full without rebuilding anything. This relies on the pickle support added in pytorch-symbolic 1.2.

PARAMETER	DESCRIPTION
path	File path to save to. Parent directories are created if needed. TYPE: str or Path
**metadata	Additional metadata stored alongside the model (e.g. training info). TYPE: Any DEFAULT: {}

Notes

The file is a regular torch.save payload, loaded back with weights_only=False — only open files you trust. Current reservoir states are captured as-is; reset or warm up beforehand to control what is persisted.

Any custom callable passed as a topology, *_initializer, or activation spec must be importable (a module-level def, not a lambda or locally-defined function) for the model to pickle. Specs given as strings, (name, kwargs) tuples, or registered objects always serialize. If a callable is not picklable, use the lighter state-dict :meth:save instead.

Examples:

>>> model.save_full("model_full.pt", epoch=10)
>>> restored = ESNModel.load_full("model_full.pt")  # no rebuild needed
>>> predictions = restored.forecast(warmup, horizon=100)

See Also

load_full : Reconstruct a model saved with this method. save : Lighter, state-dict-only persistence (architecture not stored).

Source code in src/resdag/core/model.py

def save_full(
    self,
    path: str | Path,
    **metadata: Any,
) -> None:
    """
    Serialize the *entire* model — architecture, weights, and reservoir
    states — to a single file.

    Unlike :meth:`save`, which stores only the ``state_dict`` (so the
    architecture must be re-created before :meth:`load`), this pickles the
    whole model object, the ``pytorch_symbolic`` graph included.  Restore it
    with :meth:`load_full` without rebuilding anything.  This relies on the
    pickle support added in ``pytorch-symbolic`` 1.2.

    Parameters
    ----------
    path : str or Path
        File path to save to. Parent directories are created if needed.
    **metadata
        Additional metadata stored alongside the model (e.g. training info).

    Notes
    -----
    The file is a regular ``torch.save`` payload, loaded back with
    ``weights_only=False`` — only open files you trust.  Current reservoir
    states are captured as-is; reset or warm up beforehand to control what
    is persisted.

    Any custom callable passed as a ``topology``, ``*_initializer``, or
    ``activation`` spec must be importable (a module-level ``def``, not a
    ``lambda`` or locally-defined function) for the model to pickle.  Specs
    given as strings, ``(name, kwargs)`` tuples, or registered objects
    always serialize.  If a callable is not picklable, use the lighter
    state-dict :meth:`save` instead.

    Examples
    --------
    >>> model.save_full("model_full.pt", epoch=10)
    >>> restored = ESNModel.load_full("model_full.pt")  # no rebuild needed
    >>> predictions = restored.forecast(warmup, horizon=100)

    See Also
    --------
    load_full : Reconstruct a model saved with this method.
    save : Lighter, state-dict-only persistence (architecture not stored).
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"resdag_full_model": self, "metadata": metadata}, path)

load_full classmethod

load_full(path: str | Path, return_metadata: bool = False, map_location: Any = None) -> 'ESNModel | tuple[ESNModel, dict[str, Any]]'

Reconstruct a complete model saved with :meth:save_full.

No pre-built architecture is required — the model object, its symbolic graph, weights, and reservoir states are all restored from the file.

PARAMETER	DESCRIPTION
path	File path written by :meth:save_full. TYPE: str or Path
return_metadata	If True, return (model, metadata) instead of just the model. TYPE: bool DEFAULT: False
map_location	Passed to torch.load to remap storage devices — e.g. "cpu" to load a GPU-saved model on a CPU-only machine. TYPE: optional DEFAULT: None

RETURNS	DESCRIPTION
ESNModel or tuple of (ESNModel, dict)	The reconstructed model, optionally paired with its metadata dict.

RAISES	DESCRIPTION
ValueError	If the file does not contain a whole model — e.g. it is a state-dict checkpoint from :meth:save (rebuild the architecture and use :meth:load / :meth:load_from_file for those).

Warnings

Loads with weights_only=False, which unpickles arbitrary Python objects. Only call this on files from a source you trust.

Notes

Accepts both :meth:save_full files (a metadata wrapper) and bare torch.save(model) files (the model object on its own); the latter carry no metadata.

Examples:

>>> model.save_full("model_full.pt")
>>> restored = ESNModel.load_full("model_full.pt")
>>> model_cpu = ESNModel.load_full("gpu_model.pt", map_location="cpu")

See Also

save_full : Serialize a complete model. load_from_file : Load a state-dict checkpoint into an existing model.

Source code in src/resdag/core/model.py

@classmethod
def load_full(
    cls,
    path: str | Path,
    return_metadata: bool = False,
    map_location: Any = None,
) -> "ESNModel | tuple[ESNModel, dict[str, Any]]":
    """
    Reconstruct a complete model saved with :meth:`save_full`.

    No pre-built architecture is required — the model object, its symbolic
    graph, weights, and reservoir states are all restored from the file.

    Parameters
    ----------
    path : str or Path
        File path written by :meth:`save_full`.
    return_metadata : bool, default=False
        If True, return ``(model, metadata)`` instead of just the model.
    map_location : optional
        Passed to ``torch.load`` to remap storage devices — e.g.
        ``"cpu"`` to load a GPU-saved model on a CPU-only machine.

    Returns
    -------
    ESNModel or tuple of (ESNModel, dict)
        The reconstructed model, optionally paired with its metadata dict.

    Raises
    ------
    ValueError
        If the file does not contain a whole model — e.g. it is a
        state-dict checkpoint from :meth:`save` (rebuild the architecture
        and use :meth:`load` / :meth:`load_from_file` for those).

    Warnings
    --------
    Loads with ``weights_only=False``, which unpickles arbitrary Python
    objects.  Only call this on files from a source you trust.

    Notes
    -----
    Accepts both :meth:`save_full` files (a metadata wrapper) and bare
    ``torch.save(model)`` files (the model object on its own); the latter
    carry no metadata.

    Examples
    --------
    >>> model.save_full("model_full.pt")
    >>> restored = ESNModel.load_full("model_full.pt")
    >>> model_cpu = ESNModel.load_full("gpu_model.pt", map_location="cpu")

    See Also
    --------
    save_full : Serialize a complete model.
    load_from_file : Load a state-dict checkpoint into an existing model.
    """
    path = Path(path)
    payload = torch.load(path, weights_only=False, map_location=map_location)
    model: ESNModel
    metadata: dict[str, Any]
    if isinstance(payload, dict) and "resdag_full_model" in payload:
        model = payload["resdag_full_model"]
        metadata = payload.get("metadata", {})
    elif isinstance(payload, ESNModel):
        # A bare ``torch.save(model)`` file, with no metadata wrapper.
        model = payload
        metadata = {}
    else:
        raise ValueError(
            f"{path} does not contain a full ESNModel. Write one with "
            f"save_full() (or torch.save(model)); for a state-dict checkpoint "
            f"from save(), rebuild the architecture and use load()/load_from_file()."
        )
    if return_metadata:
        return model, metadata
    return model

plot_model

plot_model(show_shapes: bool = False, show_trainable: bool = False, rankdir: str = 'TB', save_path: str | Path | None = None, format: str = 'svg', **kwargs: Any) -> Any

Visualize model architecture as a graph.

PARAMETER	DESCRIPTION
show_shapes	Show tensor shapes on edges. TYPE: bool DEFAULT: False
show_trainable	Show a padlock indicator (🔒 frozen / 🔓 trainable) on nodes that have learnable parameters. TYPE: bool DEFAULT: False
rankdir	Graph layout direction. TYPE: ('TB', 'LR', 'BT', 'RL') DEFAULT: 'TB'
save_path	Render and save to this path instead of displaying. TYPE: str or Path DEFAULT: None
format	Output format when save_path is given. TYPE: ('svg', 'png', 'pdf') DEFAULT: 'svg'

RETURNS	DESCRIPTION
Source or None	None in Jupyter (diagram already displayed). graphviz.Source in script/REPL (system viewer opened). DOT string if graphviz is not installed.

Notes

Requires the graphviz Python package and system binary. pip install graphviz and apt install graphviz.

Source code in src/resdag/core/model.py

def plot_model(
    self,
    show_shapes: bool = False,
    show_trainable: bool = False,
    rankdir: str = "TB",
    save_path: str | Path | None = None,
    format: str = "svg",
    **kwargs: Any,
) -> Any:
    """
    Visualize model architecture as a graph.

    Parameters
    ----------
    show_shapes : bool, default=False
        Show tensor shapes on edges.
    show_trainable : bool, default=False
        Show a padlock indicator (🔒 frozen / 🔓 trainable) on nodes
        that have learnable parameters.
    rankdir : {'TB', 'LR', 'BT', 'RL'}, default='TB'
        Graph layout direction.
    save_path : str or Path, optional
        Render and save to this path instead of displaying.
    format : {'svg', 'png', 'pdf'}, default='svg'
        Output format when ``save_path`` is given.

    Returns
    -------
    graphviz.Source or None
        ``None`` in Jupyter (diagram already displayed). ``graphviz.Source``
        in script/REPL (system viewer opened). DOT string if graphviz is
        not installed.

    Notes
    -----
    Requires the ``graphviz`` Python package and system binary.
    ``pip install graphviz`` and ``apt install graphviz``.
    """
    # ── Palette ───────────────────────────────────────────────────────────
    _FONT = "#1A2332"
    _EDGE = "#8B9CB6"
    _FONTS = "Helvetica Neue,Helvetica,Arial,sans-serif"

    # ── Helpers ───────────────────────────────────────────────────────────
    node_to_name: dict[Any, str] = getattr(self, "_node_to_layer_name", {})

    def _is_jupyter() -> bool:
        try:
            return bool(get_ipython().__class__.__name__ == "ZMQInteractiveShell")  # type: ignore[name-defined]
        except NameError:
            return False

    def _get_node_label(node: Any) -> str:
        if node in node_to_name:
            return node_to_name[node]
        for i, inp in enumerate(self.inputs):
            if node is inp:
                return f"Input_{i + 1}"
        return f"node_{id(node)}"

    def _get_shape_str(node: Any) -> str:
        if hasattr(node, "shape") and isinstance(node.shape, torch.Size):
            return str(tuple(node.shape))
        return ""

    def _trainable_status(module: Any) -> bool | None:
        if hasattr(module, "trainable"):
            return bool(module.trainable)
        params = list(module.parameters())
        return None if not params else any(p.requires_grad for p in params)

    def _build_dot() -> str:
        nodes: dict[str, tuple[str, str, Any, bool, bool]] = {}
        edges: list[tuple[str, str, str]] = []

        for i, inp in enumerate(self.inputs):
            nodes[f"Input_{i + 1}"] = ("Input", _get_shape_str(inp), None, True, False)

        for node, layer_name in node_to_name.items():
            module = getattr(node, "layer", None)
            cls_name = (
                type(module).__name__ if module is not None else layer_name.rsplit("_", 1)[0]
            )
            nodes[layer_name] = (cls_name, _get_shape_str(node), module, False, False)
            for parent in getattr(node, "_parents", []):
                edge_label = _get_shape_str(parent) if show_shapes else ""
                edges.append((_get_node_label(parent), layer_name, edge_label))

        for out in self.outputs:
            out_name = _get_node_label(out)
            if out_name in nodes:
                cls_name, shape_str, module, is_input, _ = nodes[out_name]
                nodes[out_name] = (cls_name, shape_str, module, is_input, True)

        # Assign hues via golden ratio over sorted unique class names.
        # Golden ratio on sequential *indices* guarantees maximal perceptual
        # separation — the only correct way to avoid adjacent colors.
        # Sorting makes assignment deterministic regardless of graph order.
        sorted_classes = sorted({cls for cls, *_ in nodes.values()})
        _hues = {cls: (i * 0.618033988749895) % 1.0 for i, cls in enumerate(sorted_classes)}

        def _hex(r: float, g: float, b: float) -> str:
            return f"#{int(r * 255):02X}{int(g * 255):02X}{int(b * 255):02X}"

        def _color_for(cls_name: str) -> tuple[str, str]:
            hue = _hues.get(cls_name, 0.0)
            fill = _hex(*colorsys.hsv_to_rgb(hue, 0.20, 0.97))
            border = _hex(*colorsys.hsv_to_rgb(hue, 0.82, 0.52))
            return fill, border

        lines = [
            "digraph ESNModel {",
            "  bgcolor=white;",
            f"  rankdir={rankdir};",
            "  splines=true;",
            "  nodesep=0.6;",
            "  ranksep=0.8;",
            "  pad=0.4;",
            f'  graph [fontname="{_FONTS}"];',
            f'  node [fontname="{_FONTS}", penwidth=1.5];',
            f'  edge [color="{_EDGE}", penwidth=1.2, arrowsize=0.7, arrowhead=vee,'
            f' fontname="{_FONTS}", fontcolor="{_EDGE}", fontsize=9];',
        ]

        for name, (cls_name, shape_str, module, is_input, is_output) in nodes.items():
            lock = ""
            if show_trainable and module is not None:
                status = _trainable_status(module)
                if status is not None:
                    lock = " \U0001f513" if status else " \U0001f512"
            label = f'"{cls_name}{lock}"'
            if show_shapes and shape_str:
                label = f'"{cls_name}{lock}\\n{shape_str}"'

            fill, border = _color_for(cls_name)
            node_shape = "ellipse" if is_input else "box"
            style = (
                "filled"
                if is_input
                else ("filled,rounded,bold" if is_output else "filled,rounded")
            )
            lines.append(
                f'  "{name}" [label={label}, shape={node_shape},'
                f' style="{style}", fillcolor="{fill}", color="{border}",'
                f' fontcolor="{_FONT}"];'
            )

        for from_name, to_name, edge_label in edges:
            if edge_label:
                lines.append(f'  "{from_name}" -> "{to_name}" [label="{edge_label}"];')
            else:
                lines.append(f'  "{from_name}" -> "{to_name}";')

        lines.append("}")
        return "\n".join(lines)

    # ── Rendering ─────────────────────────────────────────────────────────
    def _print_dot_fallback(reason: str) -> str:
        print(reason)
        print("DOT source (paste at https://dreampuf.github.io/GraphvizOnline/):")
        dot_src = _build_dot()
        print(dot_src)
        return dot_src

    try:
        try:
            import graphviz
        except ImportError as exc:
            # graphviz is an optional dependency; the Python package may be
            # absent. Re-raise with an actionable install hint (the outer
            # handler turns any ImportError into the printable DOT fallback).
            raise ImportError(
                "graphviz Python package not installed; install it with "
                "`pip install resdag[viz]` or `pip install graphviz` "
                "(the system binary is also required: e.g. apt install graphviz)."
            ) from exc

        dot_src = _build_dot()
        src = graphviz.Source(dot_src)

        if save_path is not None:
            save_path = Path(save_path)
            try:
                src.render(str(save_path.with_suffix("")), format=format, cleanup=True)
            except graphviz.ExecutableNotFound:
                return _print_dot_fallback(
                    "graphviz system binary not found: install it "
                    "(e.g. apt install graphviz) or render the DOT source below."
                )
            print(f"Saved to {save_path.with_suffix('.' + format)}")
            return src

        if _is_jupyter():
            from IPython.display import SVG
            from IPython.display import display as ipy_display

            try:
                svg = src.pipe(format="svg").decode("utf-8")
            except graphviz.ExecutableNotFound:
                return _print_dot_fallback(
                    "graphviz system binary not found: install it "
                    "(e.g. apt install graphviz) or render the DOT source below."
                )
            ipy_display(SVG(svg))
            return None  # prevent double-display from cell output

        try:
            src.view(cleanup=True)
        except Exception:
            pass
        return src

    except ImportError as exc:
        return _print_dot_fallback(
            f"{exc} pip install resdag[viz] (or pip install graphviz) "
            f"&& apt install graphviz, or render the DOT source below."
        )

warmup

warmup(*inputs: Tensor, return_outputs: bool = False, reset: bool = True, no_grad: bool = True) -> Tensor | tuple[Tensor, ...] | None

Teacher-forced warmup to synchronize reservoir states.

Runs the model forward with provided inputs, updating internal reservoir states to achieve the Echo State Property (synchronization with input dynamics).

PARAMETER	DESCRIPTION
*inputs	Input tensors of shape (batch, timesteps, features). Convention: first input is feedback, remaining are drivers. TYPE: Tensor DEFAULT: ()
return_outputs	If True, return model outputs during warmup. TYPE: bool DEFAULT: False
reset	If True, reservoir states are reset to None before the warmup pass. Set to False only if you want to continue from a previously saved state — typical workflows always reset. TYPE: bool DEFAULT: True
no_grad	If True (default), the warmup pass runs under :func:torch.no_grad — no autograd graph is built and behaviour/performance are identical to the historical @torch.no_grad() decorator. Set False to keep the graph so the teacher-forced pass is part of a wider backpropagation (e.g. BPTT over the warmup→forecast seam). For gradients to flow across the warmup→forecast boundary the reservoir must also keep its carried state attached: set reservoir.detach_state_between_calls = False (it defaults to True, which detaches the stored state at every forward call — see :class:~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer). Note that no_grad=False retains the full warmup graph in memory, so cost grows with the warmup length. TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
torch.Tensor or tuple of torch.Tensor or None	If return_outputs=True: the model output of shape (batch, timesteps, output_dim) for a single-output model, or a tuple of such tensors for a multi-output model. Otherwise: None (only internal state is updated).

RAISES	DESCRIPTION
ValueError	If no inputs are provided.

Examples:

Synchronize states without capturing output:

>>> model.warmup(feedback_data)

Synchronize and capture output:

>>> outputs = model.warmup(feedback_data, return_outputs=True)

With driving input:

>>> model.warmup(feedback, driving_signal)

Continue warming from a saved state:

>>> model.set_reservoir_states(saved_states)
>>> model.warmup(more_data, reset=False)

See Also

forecast : Two-phase forecasting with warmup and generation. reset_reservoirs : Reset all reservoir states.

Source code in src/resdag/core/model.py

def warmup(
    self,
    *inputs: torch.Tensor,
    return_outputs: bool = False,
    reset: bool = True,
    no_grad: bool = True,
) -> torch.Tensor | tuple[torch.Tensor, ...] | None:
    """
    Teacher-forced warmup to synchronize reservoir states.

    Runs the model forward with provided inputs, updating internal
    reservoir states to achieve the Echo State Property (synchronization
    with input dynamics).

    Parameters
    ----------
    *inputs : torch.Tensor
        Input tensors of shape ``(batch, timesteps, features)``.
        Convention: first input is feedback, remaining are drivers.
    return_outputs : bool, default=False
        If True, return model outputs during warmup.
    reset : bool, default=True
        If True, reservoir states are reset to ``None`` before the
        warmup pass.  Set to ``False`` only if you want to continue
        from a previously saved state — typical workflows always reset.
    no_grad : bool, default=True
        If True (default), the warmup pass runs under :func:`torch.no_grad`
        — no autograd graph is built and behaviour/performance are identical
        to the historical ``@torch.no_grad()`` decorator.  Set ``False`` to
        keep the graph so the teacher-forced pass is part of a wider
        backpropagation (e.g. BPTT over the warmup→forecast seam).

        For gradients to flow *across* the warmup→forecast boundary the
        reservoir must also keep its carried state attached: set
        ``reservoir.detach_state_between_calls = False`` (it defaults to
        ``True``, which detaches the stored state at every forward call —
        see :class:`~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer`).
        Note that ``no_grad=False`` retains the full warmup graph in memory,
        so cost grows with the warmup length.

    Returns
    -------
    torch.Tensor or tuple of torch.Tensor or None
        If ``return_outputs=True``: the model output of shape
        ``(batch, timesteps, output_dim)`` for a single-output model, or a
        tuple of such tensors for a multi-output model.
        Otherwise: None (only internal state is updated).

    Raises
    ------
    ValueError
        If no inputs are provided.

    Examples
    --------
    Synchronize states without capturing output:

    >>> model.warmup(feedback_data)

    Synchronize and capture output:

    >>> outputs = model.warmup(feedback_data, return_outputs=True)

    With driving input:

    >>> model.warmup(feedback, driving_signal)

    Continue warming from a saved state:

    >>> model.set_reservoir_states(saved_states)
    >>> model.warmup(more_data, reset=False)

    See Also
    --------
    forecast : Two-phase forecasting with warmup and generation.
    reset_reservoirs : Reset all reservoir states.
    """
    if len(inputs) == 0:
        raise ValueError("At least one input (feedback) is required")

    if reset:
        self.reset_reservoirs()

    with _grad_context(no_grad):
        output = self(*inputs)

    return output if return_outputs else None

forecast

forecast(warmup_inputs: tuple[Tensor, ...] | Tensor, forecast_inputs: tuple[Tensor, ...] | None = None, *, horizon: int, initial_feedback: Tensor | None = None, return_warmup: bool = False, reset: bool = True, compile: bool = False, no_grad: bool = True) -> Tensor | tuple[Tensor, ...]

Two-phase forecast: teacher-forced warmup + autoregressive generation.

Phase 1 (Warmup): Runs model with provided inputs to synchronize reservoir states with input dynamics (Echo State Property).

Phase 2 (Forecast): Autoregressive generation where feedback comes from the model's own output while driving inputs (if any) are supplied through forecast_inputs.

PARAMETER	DESCRIPTION
warmup_inputs	Warmup tensors of shape (batch, warmup_steps, features). Convention: first element is feedback, remaining are drivers. A single tensor is accepted for the common feedback-only case. TYPE: tuple of torch.Tensor or torch.Tensor
forecast_inputs	Driver inputs for the autoregressive phase (feedback is provided by the model's own output, so it is not part of this tuple). forecast_inputs[i][:, t, :] is the driver paired with forecast step t — pass the driver series for the forecast window, continuing exactly where the warmup drivers ended. Each tensor must provide at least horizon timesteps; only the first horizon are consumed. Required when the model has driver inputs. TYPE: tuple of torch.Tensor DEFAULT: None
horizon	Number of autoregressive steps to generate. Must be >= 1. TYPE: (int, keyword - only)
initial_feedback	Custom feedback of shape (batch, 1, feedback_dim) used to seed the first autoregressive step. If None, the last warmup output is used as the seed. TYPE: Tensor DEFAULT: None
return_warmup	If True, prepend warmup outputs to the result. TYPE: bool DEFAULT: False
reset	If True, reservoir states are reset to None before warmup. Set to False only if you want to continue from a previously saved state. TYPE: bool DEFAULT: True
compile	If True, wrap the flattened per-step engine in :func:torch.compile (mode="reduce-overhead", fullgraph=True) to fuse the autoregressive step and cut Python/launch overhead — most useful for long horizons on GPU. Requires torch>=2.10; on older torch, or if compilation fails, it transparently falls back to the eager flat step (emitting a :class:RuntimeWarning). The default eager path is already a flat, graph-free step engine, so compile=False is fast; compile=True is an opt-in extra. TYPE: bool DEFAULT: False
no_grad	If True (default), the whole two-phase forecast runs under :func:torch.no_grad — no autograd graph is built and the result is bit-for-bit and performance-wise identical to the historical @torch.no_grad() decorator. Set False to make the rollout differentiable: the warmup and every autoregressive step then build a backpropagation-through-time (BPTT) graph, so calling .backward() on a loss over the returned tensor reaches the trainable reservoir and readout parameters (and any input that requires_grad). See Notes for the memory cost and the detach_state_between_calls interaction. no_grad=False is incompatible with compile=True (the reduce-overhead/cudagraph step path does not support autograd); requesting both raises ValueError. TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
torch.Tensor or tuple of torch.Tensor	For single-output models: tensor of shape (batch, horizon, output_dim) or (batch, warmup_steps + horizon, output_dim) if return_warmup. For multi-output models: tuple of tensors with the same structure. With no_grad=False the returned tensor(s) carry grad_fn / requires_grad=True whenever any participating parameter or input requires grad.

RAISES	DESCRIPTION
ValueError	If no warmup inputs are provided, if forecast_inputs is required but missing, if dimensions don't match, or if both compile=True and no_grad=False are requested.

Notes

• Convention: first warmup element is always feedback (used for autoregression).
• For multi-output models, the first output is used as feedback.
• Feedback output dimension must match feedback input dimension.
• Every returned step is genuinely autoregressive: the loop feeds the model its own previous output (seeded by initial_feedback or the last warmup output). No teacher-forced frame is included, so horizon=1 produces one real forecast step rather than echoing the warmup output, and return_warmup=True introduces no duplicated frame at the warmup/forecast seam.

Driver time alignment. Training pairs (feedback_t, driver_t) with target feedback_{t+1}. The forecast loop preserves that pairing: step t feeds the current feedback together with forecast_inputs[:, t] — the driver for that same step — and writes the resulting prediction to slot t. The seed feedback corresponds to the first step after the warmup window, so forecast_inputs must hold the driver series beginning right after the warmup drivers (no overlap) and provide at least horizon timesteps.

Flat single-step engine. The autoregressive loop is driven by a flattened, graph-free step compiled once from the model graph (see :mod:resdag.core._flat_inference) — numerically identical to the per-step graph re-execution it replaces but without the per-step nn.Module.__call__ dispatch and reservoir sequence-loop bookkeeping. Because it calls each layer's forward (and the reservoir's :meth:~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.step_stateless) directly, module forward/pre-hooks do not fire during the autoregressive steps (the teacher-forced warmup still uses the full graph path). The teacher-forced warmup runs the standard graph forward.

Differentiable rollout (no_grad=False). By default the rollout is wrapped in :func:torch.no_grad for speed and memory. Pass no_grad=False to keep the autograd graph so gradients flow through the full BPTT unroll — the canonical building block for multi-step training losses (scheduled sampling, tuning a trainable reservoir on forecast error). Two things to keep in mind:

• Memory. The graph retains every intermediate from the warmup and all horizon steps, so peak memory grows linearly with warmup_steps + horizon. Keep the horizon modest, or truncate it, when backpropagating.
• Cross-seam gradients. For the gradient to flow from a forecast step back through the warmup, the reservoir must not detach its carried state between the warmup forward and the autoregressive loop. Set reservoir.detach_state_between_calls = False (it defaults to True; the detach happens in :meth:~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.forward). The autoregressive loop itself threads the state explicitly and never detaches, so gradients always flow within the forecast window.

Examples:

Simple feedback-only model:

>>> warmup_data = torch.randn(1, 50, 3)
>>> predictions = model.forecast(warmup_data, horizon=100)
>>> print(predictions.shape)
torch.Size([1, 100, 3])

Input-driven model:

>>> predictions = model.forecast(
...     (warmup_feedback, warmup_driver),
...     forecast_inputs=(future_driver,),
...     horizon=100,
... )

Include warmup in output:

>>> full_output = model.forecast(
...     warmup_data,
...     horizon=100,
...     return_warmup=True,
... )
>>> print(full_output.shape)  # warmup_steps + horizon
torch.Size([1, 150, 3])

Differentiable multi-step rollout — train a trainable reservoir on a forecast-error loss with BPTT. Build the reservoir trainable at construction (trainable=True), then keep the warmup→forecast seam graph by turning off cross-call detaching:

>>> reservoir = ESNLayer(reservoir_size, feedback_size, trainable=True)
>>> reservoir.detach_state_between_calls = False  # keep the seam graph
>>> # ... wire ``reservoir`` into the model and readout ...
>>> opt = torch.optim.Adam(model.parameters(), lr=1e-3)
>>> for _ in range(steps):
...     opt.zero_grad()
...     preds = model.forecast(warmup_data, horizon=2, no_grad=False)
...     loss = ((preds - target) ** 2).mean()
...     loss.backward()           # reaches reservoir + readout params
...     opt.step()

See Also

warmup : Teacher-forced warmup only. reset_reservoirs : Reset reservoir states before forecasting.

Source code in src/resdag/core/model.py

def forecast(
    self,
    warmup_inputs: tuple[torch.Tensor, ...] | torch.Tensor,
    forecast_inputs: tuple[torch.Tensor, ...] | None = None,
    *,
    horizon: int,
    initial_feedback: torch.Tensor | None = None,
    return_warmup: bool = False,
    reset: bool = True,
    compile: bool = False,
    no_grad: bool = True,
) -> torch.Tensor | tuple[torch.Tensor, ...]:
    """
    Two-phase forecast: teacher-forced warmup + autoregressive generation.

    Phase 1 (Warmup): Runs model with provided inputs to synchronize
    reservoir states with input dynamics (Echo State Property).

    Phase 2 (Forecast): Autoregressive generation where feedback comes
    from the model's own output while driving inputs (if any) are
    supplied through ``forecast_inputs``.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor or torch.Tensor
        Warmup tensors of shape ``(batch, warmup_steps, features)``.
        Convention: first element is feedback, remaining are drivers.
        A single tensor is accepted for the common feedback-only case.
    forecast_inputs : tuple of torch.Tensor, optional
        Driver inputs for the autoregressive phase (feedback is provided
        by the model's own output, so it is not part of this tuple).
        ``forecast_inputs[i][:, t, :]`` is the driver paired with forecast
        step ``t`` — pass the driver series for the forecast window,
        continuing exactly where the warmup drivers ended.  Each tensor
        must provide at least ``horizon`` timesteps; only the first
        ``horizon`` are consumed.  Required when the model has driver
        inputs.
    horizon : int, keyword-only
        Number of autoregressive steps to generate. Must be ``>= 1``.
    initial_feedback : torch.Tensor, optional
        Custom feedback of shape ``(batch, 1, feedback_dim)`` used to seed
        the first autoregressive step. If None, the last warmup output is
        used as the seed.
    return_warmup : bool, default=False
        If True, prepend warmup outputs to the result.
    reset : bool, default=True
        If True, reservoir states are reset to ``None`` before warmup.
        Set to ``False`` only if you want to continue from a previously
        saved state.
    compile : bool, default=False
        If ``True``, wrap the flattened per-step engine in
        :func:`torch.compile` (``mode="reduce-overhead"``, ``fullgraph=True``)
        to fuse the autoregressive step and cut Python/launch overhead —
        most useful for long horizons on GPU.  Requires ``torch>=2.10``;
        on older torch, or if compilation fails, it transparently falls
        back to the eager flat step (emitting a :class:`RuntimeWarning`).
        The default eager path is already a flat, graph-free step engine,
        so ``compile=False`` is fast; ``compile=True`` is an opt-in extra.
    no_grad : bool, default=True
        If True (default), the whole two-phase forecast runs under
        :func:`torch.no_grad` — no autograd graph is built and the result is
        bit-for-bit and performance-wise identical to the historical
        ``@torch.no_grad()`` decorator.  Set ``False`` to make the rollout
        **differentiable**: the warmup and every autoregressive step then
        build a backpropagation-through-time (BPTT) graph, so calling
        ``.backward()`` on a loss over the returned tensor reaches the
        trainable reservoir and readout parameters (and any input that
        ``requires_grad``).  See Notes for the memory cost and the
        ``detach_state_between_calls`` interaction.

        ``no_grad=False`` is incompatible with ``compile=True`` (the
        ``reduce-overhead``/cudagraph step path does not support autograd);
        requesting both raises ``ValueError``.

    Returns
    -------
    torch.Tensor or tuple of torch.Tensor
        For single-output models: tensor of shape
        ``(batch, horizon, output_dim)`` or
        ``(batch, warmup_steps + horizon, output_dim)`` if ``return_warmup``.

        For multi-output models: tuple of tensors with the same structure.

        With ``no_grad=False`` the returned tensor(s) carry ``grad_fn`` /
        ``requires_grad=True`` whenever any participating parameter or input
        requires grad.

    Raises
    ------
    ValueError
        If no warmup inputs are provided, if ``forecast_inputs`` is required
        but missing, if dimensions don't match, or if both ``compile=True``
        and ``no_grad=False`` are requested.

    Notes
    -----
    - Convention: first warmup element is always feedback (used for autoregression).
    - For multi-output models, the first output is used as feedback.
    - Feedback output dimension must match feedback input dimension.
    - Every returned step is genuinely autoregressive: the loop feeds the
      model its own previous output (seeded by ``initial_feedback`` or the
      last warmup output). No teacher-forced frame is included, so
      ``horizon=1`` produces one real forecast step rather than echoing the
      warmup output, and ``return_warmup=True`` introduces no duplicated
      frame at the warmup/forecast seam.

    **Driver time alignment.** Training pairs ``(feedback_t, driver_t)``
    with target ``feedback_{t+1}``.  The forecast loop preserves that
    pairing: step ``t`` feeds the current feedback together with
    ``forecast_inputs[:, t]`` — the driver for that same step — and writes
    the resulting prediction to slot ``t``.  The seed feedback corresponds
    to the first step after the warmup window, so ``forecast_inputs`` must
    hold the driver series beginning right after the warmup drivers (no
    overlap) and provide at least ``horizon`` timesteps.

    **Flat single-step engine.** The autoregressive loop is driven by a
    flattened, graph-free step compiled once from the model graph (see
    :mod:`resdag.core._flat_inference`) — numerically identical to the
    per-step graph re-execution it replaces but without the per-step
    ``nn.Module.__call__`` dispatch and reservoir sequence-loop bookkeeping.
    Because it calls each layer's ``forward`` (and the reservoir's
    :meth:`~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.step_stateless`)
    directly, **module forward/pre-hooks do not fire during the
    autoregressive steps** (the teacher-forced warmup still uses the full
    graph path).  The teacher-forced warmup runs the standard graph forward.

    **Differentiable rollout (``no_grad=False``).** By default the rollout
    is wrapped in :func:`torch.no_grad` for speed and memory.  Pass
    ``no_grad=False`` to keep the autograd graph so gradients flow through
    the full BPTT unroll — the canonical building block for multi-step
    training losses (scheduled sampling, tuning a trainable reservoir on
    forecast error).  Two things to keep in mind:

    - *Memory.* The graph retains every intermediate from the warmup and all
      ``horizon`` steps, so peak memory grows linearly with
      ``warmup_steps + horizon``.  Keep the horizon modest, or truncate it,
      when backpropagating.
    - *Cross-seam gradients.* For the gradient to flow from a forecast step
      back through the warmup, the reservoir must not detach its carried
      state between the warmup forward and the autoregressive loop.  Set
      ``reservoir.detach_state_between_calls = False`` (it defaults to
      ``True``; the detach happens in
      :meth:`~resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer.forward`).
      The autoregressive loop itself threads the state explicitly and never
      detaches, so gradients always flow *within* the forecast window.

    Examples
    --------
    Simple feedback-only model:

    >>> warmup_data = torch.randn(1, 50, 3)
    >>> predictions = model.forecast(warmup_data, horizon=100)
    >>> print(predictions.shape)
    torch.Size([1, 100, 3])

    Input-driven model:

    >>> predictions = model.forecast(
    ...     (warmup_feedback, warmup_driver),
    ...     forecast_inputs=(future_driver,),
    ...     horizon=100,
    ... )

    Include warmup in output:

    >>> full_output = model.forecast(
    ...     warmup_data,
    ...     horizon=100,
    ...     return_warmup=True,
    ... )
    >>> print(full_output.shape)  # warmup_steps + horizon
    torch.Size([1, 150, 3])

    Differentiable multi-step rollout — train a trainable reservoir on a
    forecast-error loss with BPTT.  Build the reservoir trainable at
    construction (``trainable=True``), then keep the warmup→forecast seam
    graph by turning off cross-call detaching:

    >>> reservoir = ESNLayer(reservoir_size, feedback_size, trainable=True)
    >>> reservoir.detach_state_between_calls = False  # keep the seam graph
    >>> # ... wire ``reservoir`` into the model and readout ...
    >>> opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    >>> for _ in range(steps):
    ...     opt.zero_grad()
    ...     preds = model.forecast(warmup_data, horizon=2, no_grad=False)
    ...     loss = ((preds - target) ** 2).mean()
    ...     loss.backward()           # reaches reservoir + readout params
    ...     opt.step()

    See Also
    --------
    warmup : Teacher-forced warmup only.
    reset_reservoirs : Reset reservoir states before forecasting.
    """
    # Normalise warmup_inputs into a tuple
    if isinstance(warmup_inputs, torch.Tensor):
        warmup_inputs = (warmup_inputs,)
    else:
        warmup_inputs = tuple(warmup_inputs)
    if len(warmup_inputs) == 0:
        raise ValueError("At least one warmup input (feedback) is required")

    # Validate the horizon up front so callers get a clear error instead of
    # a downstream IndexError / negative-dimension RuntimeError.
    if horizon < 1:
        raise ValueError(f"horizon must be a positive integer, got {horizon}")

    # A differentiable rollout and the cudagraph/reduce-overhead compile
    # path are mutually exclusive: the compiled step reuses static output
    # buffers and is built for inference, so it cannot back-propagate.
    if compile and not no_grad:
        raise ValueError(
            "forecast(compile=True) is incompatible with no_grad=False: the "
            "reduce-overhead/cudagraph step path does not support autograd. "
            "Use no_grad=False with compile=False for a differentiable rollout."
        )

    # Determine if model has driving inputs
    num_drivers = len(warmup_inputs) - 1
    has_drivers = num_drivers > 0

    # Validate forecast_inputs (drivers only). Every autoregressive step
    # consumes one driver value, so at least ``horizon`` are required.
    if has_drivers:
        if forecast_inputs is None:
            raise ValueError(
                f"Model has {num_drivers} driving input(s). "
                f"forecast_inputs must be provided for the autoregressive phase."
            )
        forecast_inputs = tuple(forecast_inputs)
        if len(forecast_inputs) != num_drivers:
            raise ValueError(
                f"Expected {num_drivers} forecast driver(s), got {len(forecast_inputs)}"
            )
        for i, driver in enumerate(forecast_inputs):
            if driver.shape[1] < horizon:
                raise ValueError(
                    f"forecast_inputs[{i}] has {driver.shape[1]} steps, expected at "
                    f"least {horizon} (one per autoregressive step). forecast_inputs "
                    f"must hold the driver series for the forecast window, starting "
                    f"right after the warmup window."
                )

    batch_size = warmup_inputs[0].shape[0]
    feedback_dim = warmup_inputs[0].shape[-1]
    device = warmup_inputs[0].device
    dtype = warmup_inputs[0].dtype

    # Validate the seed feedback shape up front: (batch, 1, feedback_dim).
    if initial_feedback is not None:
        expected_shape = (batch_size, 1, feedback_dim)
        if initial_feedback.dim() != 3 or tuple(initial_feedback.shape) != expected_shape:
            raise ValueError(
                f"initial_feedback must have shape {expected_shape} "
                f"(batch, 1, feedback_dim); got {tuple(initial_feedback.shape)}."
            )

    # Phase 1: Warmup (handles reset internally). With ``return_outputs=True``
    # the result is never None — a single tensor for single-output models, a
    # tuple of tensors for multi-output models. ``no_grad`` is threaded
    # through so the warmup graph is kept (or not) consistently with Phase 2.
    warmup_outputs = self.warmup(
        *warmup_inputs, return_outputs=True, reset=reset, no_grad=no_grad
    )

    # Determine output structure
    output_shape = self.output_shape
    multi_output = isinstance(output_shape, tuple) and isinstance(output_shape[0], torch.Size)

    # Validate feedback dimension
    if multi_output:
        feedback_output_dim = output_shape[0][-1]
    else:
        feedback_output_dim = output_shape[-1]

    if feedback_output_dim != feedback_dim:
        # Collect readout layer names for a helpful hint.
        readout_names = []
        for name, module in self.named_modules():
            # ReadoutLayer subclasses always have an out_features attribute.
            if hasattr(module, "out_features") and hasattr(module, "_is_fitted"):
                label = getattr(module, "_name", None) or name
                readout_names.append(f"{label} (out_features={module.out_features})")
        readouts_hint = f" Model readouts: {readout_names}." if readout_names else ""
        raise ValueError(
            f"forecast(): feedback dimension mismatch. "
            f"Feedback input has {feedback_dim} features but the model output "
            f"used as feedback has {feedback_output_dim} features. "
            f"For autoregressive forecasting the first output must match the "
            f"feedback input dimension.{readouts_hint}"
        )

    # Drivers narrowed to a concrete tuple for the loop (empty for the
    # feedback-only path). Step ``t`` consumes ``drivers[i][:, t:t+1, :]``.
    drivers: tuple[torch.Tensor, ...] = forecast_inputs if forecast_inputs is not None else ()

    # Phase 2: Autoregressive forecast, driven by the flattened single-step
    # engine instead of re-walking the symbolic graph each step. Every slot
    # is produced by feeding the model its own previous output, seeded by
    # ``initial_feedback`` (if given) or the last warmup output — no
    # teacher-forced frame is emitted, so slot 0 is a genuine forecast step
    # and the warmup/forecast seam is not duplicated. The shared
    # ``rollout`` helper (see :mod:`resdag.core._flat_inference`) owns the
    # loop body, the compile/eager engine choice, and the two output-assembly
    # strategies; here we resolve the seed and initial states, delegate, and
    # persist the advanced states.
    #
    # The whole phase runs under ``_grad_context(no_grad)``: ``torch.no_grad``
    # by default (graph-free, the historical behaviour), or ``nullcontext``
    # when ``no_grad=False`` so the per-step calls build a BPTT graph.
    with _grad_context(no_grad):
        flat = self._get_flat_step()

        # Resolve reservoir states (warmup leaves them set; init any that are
        # somehow missing). They are threaded through ``rollout`` explicitly
        # rather than mutated in place each iteration.
        states: list[torch.Tensor] = []
        for reservoir in flat.reservoir_layers:
            if reservoir.state is None:
                reservoir._maybe_init_state(batch_size, device, dtype)
            assert reservoir.state is not None  # narrows for mypy
            states.append(reservoir.state)

        # Per-output feature dimensions and the (batch, 1, feedback_dim) seed.
        # The engine threads 3-D single-step slices so every non-reservoir layer
        # sees exactly what the graph path would; only the reservoir update is
        # squeezed to 2-D internally.
        if multi_output:
            assert isinstance(warmup_outputs, tuple)  # guaranteed by multi_output
            out_dims = [out.shape[-1] for out in warmup_outputs]
            seed = (
                initial_feedback
                if initial_feedback is not None
                else warmup_outputs[0][:, -1:, :]
            )
        else:
            assert isinstance(warmup_outputs, torch.Tensor)  # guaranteed by single output
            out_dims = [warmup_outputs.shape[-1]]
            seed = (
                initial_feedback if initial_feedback is not None else warmup_outputs[:, -1:, :]
            )

        forecast_buffers, states = rollout(
            flat,
            seed,
            drivers,
            states,
            horizon,
            out_dims,
            no_grad=no_grad,
            compile=compile,
        )

        # Persist the final reservoir states so a follow-up
        # ``forecast(reset=False)`` or ``get_reservoir_states()`` sees them,
        # mirroring the in-place mutation of the legacy per-step loop. Clone
        # so a compiled step's static output buffers cannot be overwritten
        # under the stored state.
        #
        # Honour the truncated-BPTT contract on the stored state exactly as
        # ``BaseReservoirLayer.forward`` and ``set_state`` do: under
        # ``detach_state_between_calls`` (the default) a grad-bearing rollout
        # state must be detached before it is stored, otherwise the next
        # ordinary ``forward``+``backward`` would try to backprop through the
        # already-freed rollout graph ("backward through the graph a second
        # time"), and the whole rollout graph would stay pinned in memory.
        # Cross-window BPTT (``windowed_forecast(no_grad=False)``) still works
        # because it requires ``detach_state_between_calls=False``, which
        # preserves the grad_fn here so a continuation shares the rollout graph.
        for reservoir, state in zip(flat.reservoir_layers, states):
            if reservoir.detach_state_between_calls and state.grad_fn is not None:
                reservoir.state = state.detach().clone()
            else:
                reservoir.state = state.clone()

        if multi_output:
            if return_warmup:
                assert isinstance(warmup_outputs, tuple)
                return tuple(
                    torch.cat([warmup_outputs[i], forecast_buffers[i]], dim=1)
                    for i in range(len(forecast_buffers))
                )
            return tuple(forecast_buffers)

        if return_warmup:
            assert isinstance(warmup_outputs, torch.Tensor)
            return torch.cat([warmup_outputs, forecast_buffers[0]], dim=1)
        return forecast_buffers[0]

windowed_forecast

windowed_forecast(series: Tensor, *driver_series: Tensor, predict_len: int, teacher_len: int, warmup_len: int | None = None, reset: bool = True, return_mask: bool = False, no_grad: bool = True) -> Tensor | tuple[Tensor, Tensor]

Gap-filling reconstruction by alternating re-sync and free-running.

Reconstruct a long trajectory that is only observed in periodic windows: the reservoir is repeatedly re-synchronized on a short window of real data (teacher forcing) and then left to forecast autonomously across the unobserved gap. Because the Echo State Property gives the reservoir fading memory, every teacher-forced window pulls its state back onto the true trajectory, no matter how far the previous free run had drifted — so per-window error does not compound across windows (accuracy within a gap still depends on keeping predict_len short relative to the system's predictability horizon).

The reservoir state is carried across the whole pass (it is reset only once, at the start, and only when reset=True). Each cycle is one :meth:forecast call: its warmup phase is the re-sync window and its autoregressive phase is the gap.

Timeline (W = warmup_len, F = teacher_len, P = predict_len)::

[== warmup W ==][~~ gap P ~~][= teacher F =][~~ gap P ~~][= F =]...
  observed         filled       observed       filled      obs
(teacher-forced)  (forecast)  (teacher-forced) (forecast)

Observed segments are copied verbatim from series; gap segments are replaced by the model's autonomous forecast. Any trailing steps that cannot host a full re-sync window plus at least one forecast step are left as observed.

PARAMETER	DESCRIPTION
series	Ground-truth feedback series of shape (batch, T, feedback_dim). It is both the source of every teacher-forced window and the value returned on observed segments. TYPE: Tensor
*driver_series	Exogenous driver series, one per driver input of the model, each of shape (batch, T, driver_dim) spanning the full timeline. Omit for feedback-only models. Unlike :meth:forecast (which takes separate warmup and forecast_inputs driver tuples), here each driver covers the whole T-step timeline and is sliced internally per cycle. TYPE: Tensor DEFAULT: ()
predict_len	Length P of each autonomous forecast window (the gap to fill). Must be >= 1. TYPE: (int, keyword - only)
teacher_len	Length F of each teacher-forced re-synchronization window. Must be >= 1. Longer windows re-sync more reliably; shorter windows observe less of the signal. TYPE: (int, keyword - only)
warmup_len	Length W of the initial teacher-forced window. Defaults to teacher_len. TYPE: int DEFAULT: None
reset	If True, reservoir states are reset before the first window. Set False to continue from the current (already warmed) state. TYPE: bool DEFAULT: True
return_mask	If True, also return a 1-D boolean mask of length T that is True where the reconstruction is a real value and False where it is a model forecast (see Returns). TYPE: bool DEFAULT: False
no_grad	If True (default), every cycle runs under :func:torch.no_grad — graph-free and identical to the historical @torch.no_grad() decorator. Set False to make the gap-fill reconstruction differentiable: the value threaded into the reconstruction on every forecast segment carries an autograd graph, so a loss over the gaps (reconstruction[:, ~mask]) can be back-propagated through the BPTT unroll of each window into the trainable reservoir and readout parameters. The flag is threaded into the per-cycle :meth:forecast calls; the same memory cost and reservoir.detach_state_between_calls caveats documented on :meth:forecast apply, multiplied across cycles (state is carried across windows, so detaching it severs gradients window-to-window). TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
reconstruction	Reconstructed series of shape (batch, T, feedback_dim): real values on observed segments, forecasts on the gaps. TYPE: Tensor
mask	Only if return_mask=True. Boolean tensor of shape (T,), True where the step is a real value (a teacher-forced window, or the untouched trailing remainder copied verbatim from series) and False where it is a model forecast. TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If series is not 3-D, if predict_len/teacher_len/ warmup_len are not positive, if the number of driver_series does not match the model's drivers, if any driver does not span T steps, or if series is too short for even one cycle.

Notes

• The forecast steps to score are exactly reconstruction[:, ~mask] against series[:, ~mask] — the values the model had to infer without ever seeing them.
• For multi-output models the first output is fed back (as in :meth:forecast) and only that channel is reconstructed.

Examples:

Reconstruct a Lorenz trajectory observed 40 steps out of every 240, filling 200-step gaps:

>>> recon, mask = model.windowed_forecast(
...     series,                 # (1, T, 3) ground truth
...     predict_len=200,
...     teacher_len=40,
...     return_mask=True,
... )
>>> gap_rmse = ((recon[:, ~mask] - series[:, ~mask]) ** 2).mean().sqrt()

With an exogenous driver:

>>> recon = model.windowed_forecast(
...     feedback_series, driver_series,
...     predict_len=100, teacher_len=50,
... )

See Also

forecast : A single warmup + autoregressive forecast cycle. warmup : Teacher-forced state synchronization only.

Source code in src/resdag/core/model.py

def windowed_forecast(
    self,
    series: torch.Tensor,
    *driver_series: torch.Tensor,
    predict_len: int,
    teacher_len: int,
    warmup_len: int | None = None,
    reset: bool = True,
    return_mask: bool = False,
    no_grad: bool = True,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """
    Gap-filling reconstruction by alternating re-sync and free-running.

    Reconstruct a long trajectory that is only *observed* in periodic
    windows: the reservoir is repeatedly re-synchronized on a short window
    of real data (teacher forcing) and then left to forecast autonomously
    across the unobserved gap.  Because the Echo State Property gives the
    reservoir fading memory, every teacher-forced window pulls its state
    back onto the true trajectory, no matter how far the previous free run
    had drifted — so per-window error does not compound across windows
    (accuracy *within* a gap still depends on keeping ``predict_len`` short
    relative to the system's predictability horizon).

    The reservoir state is carried across the whole pass (it is reset only
    once, at the start, and only when ``reset=True``).  Each cycle is one
    :meth:`forecast` call: its warmup phase is the re-sync window and its
    autoregressive phase is the gap.

    Timeline (``W = warmup_len``, ``F = teacher_len``, ``P = predict_len``)::

        [== warmup W ==][~~ gap P ~~][= teacher F =][~~ gap P ~~][= F =]...
          observed         filled       observed       filled      obs
        (teacher-forced)  (forecast)  (teacher-forced) (forecast)

    Observed segments are copied verbatim from ``series``; gap segments are
    replaced by the model's autonomous forecast.  Any trailing steps that
    cannot host a full re-sync window plus at least one forecast step are
    left as observed.

    Parameters
    ----------
    series : torch.Tensor
        Ground-truth feedback series of shape ``(batch, T, feedback_dim)``.
        It is both the source of every teacher-forced window and the value
        returned on observed segments.
    *driver_series : torch.Tensor
        Exogenous driver series, one per driver input of the model, each of
        shape ``(batch, T, driver_dim)`` spanning the full timeline.  Omit
        for feedback-only models.  Unlike :meth:`forecast` (which takes
        *separate* warmup and ``forecast_inputs`` driver tuples), here each
        driver covers the whole ``T``-step timeline and is sliced internally
        per cycle.
    predict_len : int, keyword-only
        Length ``P`` of each autonomous forecast window (the gap to fill).
        Must be ``>= 1``.
    teacher_len : int, keyword-only
        Length ``F`` of each teacher-forced re-synchronization window.
        Must be ``>= 1``.  Longer windows re-sync more reliably; shorter
        windows observe less of the signal.
    warmup_len : int, optional
        Length ``W`` of the initial teacher-forced window.  Defaults to
        ``teacher_len``.
    reset : bool, default=True
        If True, reservoir states are reset before the first window.  Set
        False to continue from the current (already warmed) state.
    return_mask : bool, default=False
        If True, also return a 1-D boolean mask of length ``T`` that is
        True where the reconstruction is a real value and False where it is
        a model forecast (see Returns).
    no_grad : bool, default=True
        If True (default), every cycle runs under :func:`torch.no_grad` —
        graph-free and identical to the historical ``@torch.no_grad()``
        decorator.  Set ``False`` to make the gap-fill reconstruction
        **differentiable**: the value threaded into the reconstruction on
        every forecast segment carries an autograd graph, so a loss over the
        gaps (``reconstruction[:, ~mask]``) can be back-propagated through the
        BPTT unroll of each window into the trainable reservoir and readout
        parameters.  The flag is threaded into the per-cycle :meth:`forecast`
        calls; the same memory cost and
        ``reservoir.detach_state_between_calls`` caveats documented on
        :meth:`forecast` apply, multiplied across cycles (state is carried
        across windows, so detaching it severs gradients window-to-window).

    Returns
    -------
    reconstruction : torch.Tensor
        Reconstructed series of shape ``(batch, T, feedback_dim)``: real
        values on observed segments, forecasts on the gaps.
    mask : torch.Tensor
        Only if ``return_mask=True``.  Boolean tensor of shape ``(T,)``,
        True where the step is a real value (a teacher-forced window, or the
        untouched trailing remainder copied verbatim from ``series``) and
        False where it is a model forecast.

    Raises
    ------
    ValueError
        If ``series`` is not 3-D, if ``predict_len``/``teacher_len``/
        ``warmup_len`` are not positive, if the number of ``driver_series``
        does not match the model's drivers, if any driver does not span
        ``T`` steps, or if ``series`` is too short for even one cycle.

    Notes
    -----
    - The forecast steps to score are exactly ``reconstruction[:, ~mask]``
      against ``series[:, ~mask]`` — the values the model had to infer
      without ever seeing them.
    - For multi-output models the first output is fed back (as in
      :meth:`forecast`) and only that channel is reconstructed.

    Examples
    --------
    Reconstruct a Lorenz trajectory observed 40 steps out of every 240,
    filling 200-step gaps:

    >>> recon, mask = model.windowed_forecast(
    ...     series,                 # (1, T, 3) ground truth
    ...     predict_len=200,
    ...     teacher_len=40,
    ...     return_mask=True,
    ... )
    >>> gap_rmse = ((recon[:, ~mask] - series[:, ~mask]) ** 2).mean().sqrt()

    With an exogenous driver:

    >>> recon = model.windowed_forecast(
    ...     feedback_series, driver_series,
    ...     predict_len=100, teacher_len=50,
    ... )

    See Also
    --------
    forecast : A single warmup + autoregressive forecast cycle.
    warmup : Teacher-forced state synchronization only.
    """
    if series.dim() != 3:
        raise ValueError(
            f"series must be 3-D (batch, T, feedback_dim); got shape {tuple(series.shape)}."
        )
    if predict_len < 1:
        raise ValueError(f"predict_len must be a positive integer, got {predict_len}")
    if teacher_len < 1:
        raise ValueError(f"teacher_len must be a positive integer, got {teacher_len}")
    if warmup_len is None:
        warmup_len = teacher_len
    if warmup_len < 1:
        raise ValueError(f"warmup_len must be a positive integer, got {warmup_len}")

    num_drivers = len(driver_series)
    expected_drivers = len(self.inputs) - 1
    if num_drivers != expected_drivers:
        raise ValueError(
            f"model has {expected_drivers} driver input(s) but {num_drivers} "
            f"driver_series were given. Pass one full-timeline driver series per "
            f"driver input (the feedback is supplied by `series`)."
        )

    _, total_steps, _ = series.shape
    for i, driver in enumerate(driver_series):
        if driver.dim() != 3:
            raise ValueError(
                f"driver_series[{i}] must be 3-D (batch, T, driver_dim); "
                f"got shape {tuple(driver.shape)}."
            )
        if driver.shape[1] != total_steps:
            raise ValueError(
                f"driver_series[{i}] spans {driver.shape[1]} steps but series spans "
                f"{total_steps}; drivers must cover the full timeline."
            )

    # ``no_grad`` is also threaded into each per-cycle ``forecast`` below; the
    # context here additionally covers the in-place reconstruction writes so
    # that, when differentiable, the forecast segments graft their autograd
    # graph onto the (cloned) reconstruction via ``CopySlices``.
    reconstruction = series.clone()
    mask = torch.ones(total_steps, dtype=torch.bool, device=series.device)

    pos = 0
    window_len = warmup_len
    ran_a_cycle = False
    with _grad_context(no_grad):
        while True:
            teacher_end = pos + window_len
            # Stop once a full re-sync window leaves no room to forecast. The
            # guard guarantees ``teacher_end < total_steps`` below, so the gap
            # (clamped to whatever remains) is always at least one step.
            if teacher_end >= total_steps:
                break
            gap = min(predict_len, total_steps - teacher_end)

            teacher_window = series[:, pos:teacher_end, :]
            if num_drivers:
                warmup_inputs: tuple[torch.Tensor, ...] = (
                    teacher_window,
                    *(d[:, pos:teacher_end, :] for d in driver_series),
                )
                forecast_inputs: tuple[torch.Tensor, ...] | None = tuple(
                    d[:, teacher_end : teacher_end + gap, :] for d in driver_series
                )
            else:
                warmup_inputs = (teacher_window,)
                forecast_inputs = None

            # ``return_warmup=True`` prepends the teacher-forced window's
            # outputs so we can recover the *last* warmup output — the model's
            # prediction of ``series[teacher_end]``, the first step of the gap.
            # By the forecast seam convention (slot ``t`` is the prediction of
            # index ``teacher_end + 1 + t``), the raw autoregressive block is
            # shifted one step ahead of the gap it must fill; without the last
            # warmup output every gap value would land one dt early. See the
            # slice below.
            full = self.forecast(
                warmup_inputs,
                forecast_inputs=forecast_inputs,
                horizon=gap,
                return_warmup=True,
                reset=reset and not ran_a_cycle,
                no_grad=no_grad,
            )
            feedback_full = full[0] if isinstance(full, tuple) else full
            # ``full`` spans ``[window_len warmup outputs][gap AR steps]``.
            # The gap must be filled with the predictions of
            # ``series[teacher_end .. teacher_end + gap - 1]``:
            #   * ``full[window_len - 1]``  -> prediction of ``series[teacher_end]``
            #     (last warmup output), and
            #   * ``full[window_len .. window_len + gap - 2]`` -> predictions of
            #     ``series[teacher_end + 1 .. teacher_end + gap - 1]``.
            # i.e. ``full[window_len - 1 : window_len + gap - 1]``. The final AR
            # slot (prediction of one step past the gap) is discarded.
            gap_preds = feedback_full[:, window_len - 1 : window_len + gap - 1, :]
            reconstruction[:, teacher_end : teacher_end + gap, :] = gap_preds
            mask[teacher_end : teacher_end + gap] = False

            ran_a_cycle = True
            pos = teacher_end + gap
            window_len = teacher_len

    if not ran_a_cycle:
        raise ValueError(
            f"series length ({total_steps}) is too short for even one cycle: the "
            f"initial warmup window (warmup_len={warmup_len}) leaves no room to "
            f"forecast (need at least warmup_len + 1 = {warmup_len + 1} steps; the "
            f"final gap is clamped to whatever remains)."
        )

    if return_mask:
        return reconstruction, mask
    return reconstruction

ReservoirFeatureExtractor

ReservoirFeatureExtractor(reservoir_size: int | Iterable[int] | None = None, feedback_size: int | None = None, input_size: int | None = None, trainable: bool = False, layers: BaseReservoirLayer | Iterable[BaseReservoirLayer] | None = None, **layer_kwargs: object)

Bases: Module

nn.Sequential-friendly reservoir feature extractor.

Wraps one reservoir layer (or a stack of them) as a plain :class:torch.nn.Module that maps a feedback sequence (batch, timesteps, feedback_size) (plus an optional driving input) to reservoir features (batch, timesteps, reservoir_size). Because the feedback-only case takes a single positional input, the extractor drops straight into a :class:torch.nn.Sequential ahead of any trainable head::

model = nn.Sequential(
    ReservoirFeatureExtractor(reservoir_size=300, feedback_size=3),
    nn.Linear(300, 3),
)

The reservoir is frozen by default (its parameters do not require gradients), so a single optimizer over model.parameters() trains only the head — the canonical reservoir-computing setup. Pass trainable=True (or call :meth:unfreeze) to backpropagate through the recurrence as well.

Stacking is supported: passing several feedback dimensions, or constructing the extractor from pre-built layers, chains the reservoirs so each consumes the previous reservoir's features as its feedback signal. A driving input, when supplied, is passed to the first reservoir only.

PARAMETER	DESCRIPTION
reservoir_size	Reservoir width. An int builds a single reservoir; a sequence of int builds a stack (one reservoir per entry, in order). Required unless layers is given. TYPE: int or sequence of int DEFAULT: None
feedback_size	Dimension of the feedback signal entering the first reservoir. Required unless layers is given. For stacked reservoirs only the first layer's feedback size is taken from this argument; the feedback size of each subsequent layer is the previous layer's reservoir size. TYPE: int DEFAULT: None
input_size	Dimension of an optional driving input fed to the first reservoir. If None (default), the extractor operates in the feedback-only, single-positional-input mode required by :class:torch.nn.Sequential. TYPE: int DEFAULT: None
trainable	If True, the reservoir parameters require gradients (full BPTT through the recurrence). The default False freezes them, the standard reservoir-computing configuration. TYPE: bool DEFAULT: False
layers	Pre-built reservoir layer(s) to wrap directly instead of constructing new ones. Any :class:~resdag.layers.reservoirs.BaseReservoirLayer subclass is accepted (ESNLayer, NGReservoirLayer, future reservoir types), not just :class:ESNLayer. When given, reservoir_size / feedback_size / input_size must be omitted (the layers already define them), but trainable still applies as a freeze/unfreeze toggle over the provided layers. The layers are wrapped by reference — their parameters are shared, not copied. TYPE: BaseReservoirLayer or sequence of BaseReservoirLayer DEFAULT: None
**layer_kwargs	Extra keyword arguments forwarded to every :class:ESNLayer that this constructor builds (for example spectral_radius, leak_rate, activation, topology, seed). Ignored when layers is provided. TYPE: object DEFAULT: {}

ATTRIBUTE	DESCRIPTION
reservoirs	The wrapped reservoir layer(s), in feed order. TYPE: ModuleList
output_size	Feature dimension of the extractor's output (the last reservoir's reservoir_size). TYPE: int

RAISES	DESCRIPTION
ValueError	If neither layers nor (reservoir_size, feedback_size) is given, if both forms are mixed, or if layers is empty.

Examples:

Drop straight into a :class:torch.nn.Sequential and train the head with a single optimizer:

>>> import torch
>>> import torch.nn as nn
>>> from resdag import ReservoirFeatureExtractor
>>> model = nn.Sequential(
...     ReservoirFeatureExtractor(reservoir_size=64, feedback_size=3),
...     nn.Linear(64, 3),
... )
>>> x = torch.randn(2, 50, 3)  # (batch, time, features)
>>> y = model(x)
>>> y.shape
torch.Size([2, 50, 3])

The reservoir is frozen by default, so only the head receives gradients:

>>> extractor = model[0]
>>> all(not p.requires_grad for p in extractor.parameters())
True

Re-zero the stateful reservoir between epochs:

>>> for epoch in range(3):
...     extractor.on_epoch_start()  # alias of reset_state()
...     # ... train one epoch ...

Reuse the reservoir of an existing :class:~resdag.core.ESNModel (shared, not copied):

>>> from resdag import ESNModel, ESNLayer, reservoir_input
>>> inp = reservoir_input(3)
>>> states = ESNLayer(64, feedback_size=3)(inp)
>>> esn = ESNModel(inp, states)
>>> extractor = ReservoirFeatureExtractor.from_model(esn)

See Also

resdag.layers.reservoirs.ESNLayer : The reservoir layer being wrapped. resdag.core.ESNModel : Build an extractor from an existing model.

Source code in src/resdag/core/feature_extractor.py

def __init__(
    self,
    reservoir_size: int | Iterable[int] | None = None,
    feedback_size: int | None = None,
    input_size: int | None = None,
    trainable: bool = False,
    layers: BaseReservoirLayer | Iterable[BaseReservoirLayer] | None = None,
    **layer_kwargs: object,
) -> None:
    super().__init__()

    if layers is not None:
        if reservoir_size is not None or feedback_size is not None or input_size is not None:
            raise ValueError(
                "Pass either pre-built `layers` or "
                "`reservoir_size`/`feedback_size`/`input_size`, not both."
            )
        if layer_kwargs:
            raise ValueError(
                "Extra layer keyword arguments are not supported when "
                "wrapping pre-built `layers`."
            )
        built = self._wrap_layers(layers)
    else:
        built = self._build_layers(
            reservoir_size=reservoir_size,
            feedback_size=feedback_size,
            input_size=input_size,
            **layer_kwargs,
        )

    self.reservoirs: nn.ModuleList = nn.ModuleList(built)
    # Apply the freeze/unfreeze toggle uniformly over the wrapped layers.
    self.unfreeze() if trainable else self.freeze()

is_frozen property

is_frozen: bool

bool: True when no wrapped reservoir parameter requires gradients.

output_size property

output_size: int

int: Feature dimension of the extractor output (last reservoir width).

from_model classmethod

from_model(esn_model: Module, trainable: bool = False) -> ReservoirFeatureExtractor

Build an extractor that reuses the reservoir layers of an existing model.

The reservoir layers of esn_model are wrapped by reference — the returned extractor shares their parameters (and stateful buffers) with the source model, it does not copy them. Training (or freezing) one therefore affects the other.

PARAMETER	DESCRIPTION
esn_model	A model containing one or more :class:~resdag.layers.reservoirs.BaseReservoirLayer submodules (typically an :class:~resdag.core.ESNModel). Reservoirs are collected in :func:~torch.nn.Module.named_modules order. TYPE: Module
trainable	Freeze/unfreeze toggle applied to the shared reservoir layers after wrapping. Because the layers are shared, this also changes requires_grad on the source model's reservoir parameters. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
ReservoirFeatureExtractor	Extractor wrapping the model's reservoir layers by reference.

RAISES	DESCRIPTION
ValueError	If esn_model contains no reservoir layers.

Examples:

>>> from resdag import ESNModel, ESNLayer, reservoir_input
>>> from resdag import ReservoirFeatureExtractor
>>> inp = reservoir_input(3)
>>> states = ESNLayer(64, feedback_size=3)(inp)
>>> esn = ESNModel(inp, states)
>>> extractor = ReservoirFeatureExtractor.from_model(esn)
>>> src = next(m for m in esn.modules() if hasattr(m, "weight_hh"))
>>> extractor.reservoirs[0].weight_hh is src.weight_hh
True

Source code in src/resdag/core/feature_extractor.py

@classmethod
def from_model(
    cls,
    esn_model: nn.Module,
    trainable: bool = False,
) -> ReservoirFeatureExtractor:
    """
    Build an extractor that reuses the reservoir layers of an existing model.

    The reservoir layers of ``esn_model`` are wrapped **by reference** — the
    returned extractor shares their parameters (and stateful buffers) with
    the source model, it does not copy them.  Training (or freezing) one
    therefore affects the other.

    Parameters
    ----------
    esn_model : torch.nn.Module
        A model containing one or more
        :class:`~resdag.layers.reservoirs.BaseReservoirLayer` submodules
        (typically an :class:`~resdag.core.ESNModel`).  Reservoirs are
        collected in :func:`~torch.nn.Module.named_modules` order.
    trainable : bool, default=False
        Freeze/unfreeze toggle applied to the shared reservoir layers after
        wrapping.  Because the layers are shared, this also changes
        ``requires_grad`` on the source model's reservoir parameters.

    Returns
    -------
    ReservoirFeatureExtractor
        Extractor wrapping the model's reservoir layers by reference.

    Raises
    ------
    ValueError
        If ``esn_model`` contains no reservoir layers.

    Examples
    --------
    >>> from resdag import ESNModel, ESNLayer, reservoir_input
    >>> from resdag import ReservoirFeatureExtractor
    >>> inp = reservoir_input(3)
    >>> states = ESNLayer(64, feedback_size=3)(inp)
    >>> esn = ESNModel(inp, states)
    >>> extractor = ReservoirFeatureExtractor.from_model(esn)
    >>> src = next(m for m in esn.modules() if hasattr(m, "weight_hh"))
    >>> extractor.reservoirs[0].weight_hh is src.weight_hh
    True
    """
    reservoirs = [
        module
        for _, module in esn_model.named_modules()
        if isinstance(module, BaseReservoirLayer)
    ]
    if not reservoirs:
        raise ValueError(
            "`esn_model` contains no reservoir layers (BaseReservoirLayer) "
            "to reuse. Build the extractor from sizes instead."
        )
    return cls(layers=reservoirs, trainable=trainable)

forward

forward(feedback: Tensor, *driving_inputs: Tensor) -> Tensor

Map a feedback sequence to reservoir features.

PARAMETER	DESCRIPTION
feedback	Feedback signal of shape (batch, timesteps, feedback_size). TYPE: Tensor
*driving_inputs	Optional driving input of shape (batch, timesteps, input_size), passed to the first reservoir only. At most one is supported. TYPE: Tensor DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	Reservoir features of shape (batch, timesteps, output_size).

RAISES	DESCRIPTION
ValueError	If more than one driving input is supplied.

Notes

The wrapped reservoirs are stateful: state carries across forward calls until :meth:reset_state (or its :meth:on_epoch_start alias) is invoked. When the extractor sits inside a :class:torch.nn.Sequential only feedback is passed, satisfying the single-positional-input contract.

Source code in src/resdag/core/feature_extractor.py

def forward(
    self,
    feedback: torch.Tensor,
    *driving_inputs: torch.Tensor,
) -> torch.Tensor:
    """
    Map a feedback sequence to reservoir features.

    Parameters
    ----------
    feedback : torch.Tensor
        Feedback signal of shape ``(batch, timesteps, feedback_size)``.
    *driving_inputs : torch.Tensor
        Optional driving input of shape ``(batch, timesteps, input_size)``,
        passed to the **first** reservoir only.  At most one is supported.

    Returns
    -------
    torch.Tensor
        Reservoir features of shape ``(batch, timesteps, output_size)``.

    Raises
    ------
    ValueError
        If more than one driving input is supplied.

    Notes
    -----
    The wrapped reservoirs are stateful: state carries across forward calls
    until :meth:`reset_state` (or its :meth:`on_epoch_start` alias) is
    invoked.  When the extractor sits inside a :class:`torch.nn.Sequential`
    only ``feedback`` is passed, satisfying the single-positional-input
    contract.
    """
    if len(driving_inputs) > 1:
        raise ValueError("Only one driving input tensor allowed")

    # The driving input enters the first reservoir; deeper reservoirs
    # consume the previous layer's features as their (sole) feedback.
    layers = self._layers
    features: torch.Tensor = layers[0](feedback, *driving_inputs)
    for reservoir in layers[1:]:
        features = reservoir(features)
    return features

reset_state

reset_state(batch_size: int | None = None) -> None

Reset every wrapped reservoir's internal state.

PARAMETER	DESCRIPTION
batch_size	If given, materialise a zero state with this batch size on each reservoir's device/dtype. If None (default), states are set to None and lazily re-initialised on the next forward pass. TYPE: int DEFAULT: None

See Also

on_epoch_start : Alias intended for use as an epoch-reset hook.

Source code in src/resdag/core/feature_extractor.py

def reset_state(self, batch_size: int | None = None) -> None:
    """
    Reset every wrapped reservoir's internal state.

    Parameters
    ----------
    batch_size : int, optional
        If given, materialise a zero state with this batch size on each
        reservoir's device/dtype.  If ``None`` (default), states are set to
        ``None`` and lazily re-initialised on the next forward pass.

    See Also
    --------
    on_epoch_start : Alias intended for use as an epoch-reset hook.
    """
    for reservoir in self._layers:
        reservoir.reset_state(batch_size=batch_size)

on_epoch_start

on_epoch_start(batch_size: int | None = None) -> None

Epoch-reset hook: re-zero the stateful reservoir between epochs.

A thin, intent-revealing alias of :meth:reset_state meant to be called at the top of each training epoch so a trajectory left over from the previous epoch never bleeds into the next one.

PARAMETER	DESCRIPTION
batch_size	Forwarded to :meth:reset_state. TYPE: int DEFAULT: None

Examples:

>>> for epoch in range(num_epochs):
...     extractor.on_epoch_start()
...     for batch in loader:
...         ...  # train the head

Source code in src/resdag/core/feature_extractor.py

def on_epoch_start(self, batch_size: int | None = None) -> None:
    """
    Epoch-reset hook: re-zero the stateful reservoir between epochs.

    A thin, intent-revealing alias of :meth:`reset_state` meant to be
    called at the top of each training epoch so a trajectory left over from
    the previous epoch never bleeds into the next one.

    Parameters
    ----------
    batch_size : int, optional
        Forwarded to :meth:`reset_state`.

    Examples
    --------
    >>> for epoch in range(num_epochs):  # doctest: +SKIP
    ...     extractor.on_epoch_start()
    ...     for batch in loader:
    ...         ...  # train the head
    """
    self.reset_state(batch_size=batch_size)

freeze

freeze() -> ReservoirFeatureExtractor

Freeze the reservoir: set requires_grad=False on all its parameters.

RETURNS	DESCRIPTION
ReservoirFeatureExtractor	self, to allow chaining.

Source code in src/resdag/core/feature_extractor.py

def freeze(self) -> ReservoirFeatureExtractor:
    """
    Freeze the reservoir: set ``requires_grad=False`` on all its parameters.

    Returns
    -------
    ReservoirFeatureExtractor
        ``self``, to allow chaining.
    """
    for param in self.reservoirs.parameters():
        param.requires_grad_(False)
    return self

unfreeze

unfreeze() -> ReservoirFeatureExtractor

Unfreeze the reservoir: set requires_grad=True on all its parameters.

RETURNS	DESCRIPTION
ReservoirFeatureExtractor	self, to allow chaining.

Source code in src/resdag/core/feature_extractor.py

def unfreeze(self) -> ReservoirFeatureExtractor:
    """
    Unfreeze the reservoir: set ``requires_grad=True`` on all its parameters.

    Returns
    -------
    ReservoirFeatureExtractor
        ``self``, to allow chaining.
    """
    for param in self.reservoirs.parameters():
        param.requires_grad_(True)
    return self

extra_repr

extra_repr() -> str

Return a one-line summary of the wrapped reservoir stack.

Source code in src/resdag/core/feature_extractor.py

def extra_repr(self) -> str:
    """Return a one-line summary of the wrapped reservoir stack."""
    widths = ", ".join(str(int(r.cell.output_size)) for r in self._layers)
    return f"reservoirs=[{widths}], frozen={self.is_frozen}"

reservoir_input

reservoir_input(feature_size: int, dtype: dtype | None = None) -> SymbolicData

Build a symbolic input tensor for a reservoir model.

Constructs an :func:pytorch_symbolic.Input of shape (1, feature_size) where the first axis is a placeholder for the time dimension (any sequence length is accepted at call time).

PARAMETER	DESCRIPTION
feature_size	Number of features per timestep. TYPE: int
dtype	Tensor dtype. Defaults to torch.get_default_dtype(). TYPE: dtype DEFAULT: None

RETURNS	DESCRIPTION
SymbolicTensor	Placeholder input with shape (1, feature_size).

Examples:

>>> inp = reservoir_input(3)
>>> inp.shape
torch.Size([1, 1, 3])

Source code in src/resdag/core/__init__.py

def reservoir_input(feature_size: int, dtype: torch.dtype | None = None) -> SymbolicData:
    """
    Build a symbolic input tensor for a reservoir model.

    Constructs an :func:`pytorch_symbolic.Input` of shape ``(1, feature_size)``
    where the first axis is a placeholder for the time dimension (any
    sequence length is accepted at call time).

    Parameters
    ----------
    feature_size : int
        Number of features per timestep.
    dtype : torch.dtype, optional
        Tensor dtype. Defaults to ``torch.get_default_dtype()``.

    Returns
    -------
    pytorch_symbolic.SymbolicTensor
        Placeholder input with shape ``(1, feature_size)``.

    Examples
    --------
    >>> inp = reservoir_input(3)
    >>> inp.shape
    torch.Size([1, 1, 3])
    """
    if dtype is None:
        dtype = torch.get_default_dtype()
    return ps.Input((1, feature_size), dtype=dtype)