Reference

Layers

The nn.Module components used to build models: single-step reservoir cells, the stateful sequence layers that wrap them, readout layers, and transform layers. All are importable directly from resdag.layers.

cells

ReservoirCell

Bases: Module, ABC

Abstract base for a single-timestep reservoir state update.

Owns all trainable (or frozen) parameters. Sequence iteration is handled by the enclosing :class:BaseReservoirLayer, not by the cell itself.

Notes

Subclasses must implement :attr:state_size, :attr:output_size, :meth:init_state, and :meth:forward.

inputs[0] passed to :meth:forward is always the feedback slice. Additional elements are driving inputs in the order they were passed to the layer's forward.

For cells where the output and the state are the same tensor (e.g. :class:ESNCell), output_size == state_size. For cells where they differ (e.g. :class:NGCell whose output is a feature vector but whose state is a delay buffer), the two properties return different values.

See Also

resdag.layers.esn.ESNCell : Concrete ESN cell implementation. resdag.layers.cells.ngrc_cell.NGCell : NG-RC cell implementation. resdag.layers.base.BaseReservoirLayer : Layer that drives the cell.

state_size abstractmethod property

state_size: int

Size of the state tensor (second dimension).

output_size abstractmethod property

output_size: int

Dimensionality of the per-step output vector.

init_state abstractmethod

init_state(batch_size: int, device: device, dtype: dtype) -> Tensor

Return a zero initial state tensor.

PARAMETER	DESCRIPTION
batch_size	Number of samples in the batch. TYPE: int
device	Target device. TYPE: device
dtype	Target dtype. TYPE: dtype

RETURNS	DESCRIPTION
Tensor	Zero-filled initial state.

Source code in src/resdag/layers/cells/base_cell.py

@abstractmethod
def init_state(
    self,
    batch_size: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """
    Return a zero initial state tensor.

    Parameters
    ----------
    batch_size : int
        Number of samples in the batch.
    device : torch.device
        Target device.
    dtype : torch.dtype
        Target dtype.

    Returns
    -------
    torch.Tensor
        Zero-filled initial state.
    """
    ...

forward abstractmethod

forward(inputs: list[Tensor], state: Tensor) -> tuple[Tensor, Tensor]

Compute the per-step output and next state from current inputs and state.

PARAMETER	DESCRIPTION
inputs	Per-timestep input slices, one per input stream, each of shape (batch, feature_dim). inputs[0] is always the feedback slice; additional elements are driving inputs. TYPE: list[Tensor]
state	Current state tensor. TYPE: Tensor

RETURNS	DESCRIPTION
output	Per-step output of shape (batch, output_size). TYPE: Tensor
new_state	Updated state tensor. TYPE: Tensor

Source code in src/resdag/layers/cells/base_cell.py

@abstractmethod
def forward(
    self,
    inputs: list[torch.Tensor],
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute the per-step output and next state from current inputs and state.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Per-timestep input slices, one per input stream, each of shape
        ``(batch, feature_dim)``.  ``inputs[0]`` is always the feedback
        slice; additional elements are driving inputs.
    state : torch.Tensor
        Current state tensor.

    Returns
    -------
    output : torch.Tensor
        Per-step output of shape ``(batch, output_size)``.
    new_state : torch.Tensor
        Updated state tensor.
    """
    ...

project_inputs

project_inputs(inputs: list[Tensor]) -> Tensor | None

Optional sequence-level fast path: precompute input contributions.

Cells whose pre-activation splits into an input-dependent part and a state-dependent part (e.g. the leaky ESN) can compute the input-dependent part for the whole sequence in one batched matmul, leaving only the recurrent term inside the time loop. This cuts the per-step kernel count roughly in half — the difference between the GPU being slower and faster than the CPU for typical reservoir sizes.

PARAMETER	DESCRIPTION
inputs	Full input sequences, one per stream, each of shape (batch, timesteps, feature_dim). inputs[0] is the feedback sequence. TYPE: list[Tensor]

RETURNS	DESCRIPTION
Tensor or None	Precomputed projection of shape (batch, timesteps, state_size) to be consumed step-by-step via :meth:step, or None if the cell does not support the fast path (the layer then falls back to calling :meth:forward per timestep).

Source code in src/resdag/layers/cells/base_cell.py

def project_inputs(self, inputs: list[torch.Tensor]) -> torch.Tensor | None:
    """
    Optional sequence-level fast path: precompute input contributions.

    Cells whose pre-activation splits into an input-dependent part and a
    state-dependent part (e.g. the leaky ESN) can compute the
    input-dependent part for the *whole* sequence in one batched matmul,
    leaving only the recurrent term inside the time loop.  This cuts the
    per-step kernel count roughly in half — the difference between the
    GPU being slower and faster than the CPU for typical reservoir sizes.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Full input sequences, one per stream, each of shape
        ``(batch, timesteps, feature_dim)``.  ``inputs[0]`` is the
        feedback sequence.

    Returns
    -------
    torch.Tensor or None
        Precomputed projection of shape ``(batch, timesteps, state_size)``
        to be consumed step-by-step via :meth:`step`, or ``None`` if the
        cell does not support the fast path (the layer then falls back to
        calling :meth:`forward` per timestep).
    """
    return None

step

step(projected_t: Tensor, state: Tensor) -> tuple[Tensor, Tensor]

Single-step update consuming a slice of :meth:project_inputs.

Only called by the layer when :meth:project_inputs returned a tensor. Subclasses implementing the fast path must override both methods together.

PARAMETER	DESCRIPTION
projected_t	Per-timestep slice of the precomputed projection, shape (batch, state_size). TYPE: Tensor
state	Current state tensor. TYPE: Tensor

RETURNS	DESCRIPTION
output	Per-step output. TYPE: Tensor
new_state	Updated state tensor. TYPE: Tensor

Source code in src/resdag/layers/cells/base_cell.py

def step(
    self,
    projected_t: torch.Tensor,
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Single-step update consuming a slice of :meth:`project_inputs`.

    Only called by the layer when :meth:`project_inputs` returned a
    tensor.  Subclasses implementing the fast path must override both
    methods together.

    Parameters
    ----------
    projected_t : torch.Tensor
        Per-timestep slice of the precomputed projection, shape
        ``(batch, state_size)``.
    state : torch.Tensor
        Current state tensor.

    Returns
    -------
    output : torch.Tensor
        Per-step output.
    new_state : torch.Tensor
        Updated state tensor.
    """
    raise NotImplementedError(
        f"{type(self).__name__} returned a projection from project_inputs() "
        f"but does not implement step()."
    )

validate_state

validate_state(state: Tensor) -> None

Validate that state matches the layout this cell expects.

The base implementation enforces the 2-D (batch, state_size) contract used by classical RNN-style cells (e.g. :class:ESNCell). Cells with a different state layout — for example :class:~resdag.layers.cells.ngrc_cell.NGCell whose state is a 3-D delay buffer — override this method to check their own shape.

PARAMETER	DESCRIPTION
state	Candidate state tensor to validate. TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If state does not match the expected layout. The error message names the cell class and the offending shape so callers can act on it without digging into the layer/cell internals.

Source code in src/resdag/layers/cells/base_cell.py

def validate_state(self, state: torch.Tensor) -> None:
    """
    Validate that ``state`` matches the layout this cell expects.

    The base implementation enforces the 2-D ``(batch, state_size)``
    contract used by classical RNN-style cells (e.g. :class:`ESNCell`).
    Cells with a different state layout — for example
    :class:`~resdag.layers.cells.ngrc_cell.NGCell` whose state is a 3-D
    delay buffer — override this method to check their own shape.

    Parameters
    ----------
    state : torch.Tensor
        Candidate state tensor to validate.

    Raises
    ------
    ValueError
        If ``state`` does not match the expected layout.  The error
        message names the cell class and the offending shape so callers
        can act on it without digging into the layer/cell internals.
    """
    if state.dim() != 2 or state.shape[-1] != self.state_size:
        raise ValueError(
            f"{type(self).__name__}.validate_state: expected a 2-D state of shape "
            f"(batch, {self.state_size}); got tensor of shape {tuple(state.shape)}."
        )

ESNCell

ESNCell(reservoir_size: int, feedback_size: int, input_size: int | None = None, spectral_radius: float | None = None, bias: bool = True, bias_scaling: float = 1.0, activation: str = 'tanh', leak_rate: float | Tensor | Sequence[float] = 1.0, noise: float = 0.0, trainable: bool = False, feedback_initializer: InitializerSpec = None, input_initializer: InitializerSpec = None, topology: TopologySpec = None, seed: SeedLike = None)

Bases: ReservoirCell

Single-timestep leaky Echo State Network update.

Owns all weight matrices and bias. Sequence iteration is delegated to the enclosing :class:ESNLayer.

The state update follows the standard leaky-integrator ESN equation (Jaeger 2001; Lukoševičius 2012):

.. math::

h_t = (1 - \alpha)\,h_{t-1} + \alpha\,f(W_{fb}\,x_{fb,t}
       + W_{in}\,x_{in,t} + W_{rec}\,h_{t-1} + b)

where :math:f is the activation function, :math:\alpha is the leak rate, :math:W_{fb} is the feedback weight matrix, :math:W_{in} is the (optional) input weight matrix, :math:W_{rec} is the recurrent weight matrix, and :math:b is a fixed random bias drawn from :math:\mathcal{U}(-\beta, \beta) with :math:\beta = bias_scaling.

The bias breaks the odd symmetry of tanh dynamics: without it, negated inputs produce exactly negated states, which constrains the representations the readout can draw from.

PARAMETER	DESCRIPTION
reservoir_size	Number of reservoir units (hidden state dimension). TYPE: int
feedback_size	Dimension of feedback signal. Required for all ESN cells. TYPE: int
input_size	Dimension of driving inputs. If None, no driving input weight matrix is created. TYPE: int DEFAULT: None
spectral_radius	Target spectral radius for recurrent weights. If None, no spectral radius scaling is applied. TYPE: float DEFAULT: None
bias	Whether to include a bias term. TYPE: bool DEFAULT: True
bias_scaling	Scale of the random bias: entries are drawn from uniform(-bias_scaling, bias_scaling), matching the default initialization of the feedback/input weights. Ignored when bias=False. Set to 0.0 to keep a zero bias (the historical pre-0.5 behaviour, where bias=True was effectively a no-op because the bias was zero-initialized and frozen). TYPE: float DEFAULT: 1.0
activation	Activation function for reservoir dynamics. TYPE: (tanh, relu, identity, sigmoid) DEFAULT: 'tanh'
leak_rate	Leaky integration rate in (0, 1]. A scalar value of 1.0 means no leaking (standard RNN update); smaller values create slower dynamics. Values outside (0, 1] are rejected at construction with a ValueError: 0.0 would freeze the state entirely and values > 1.0 (or < 0) produce ill-defined dynamics. A per-neuron leak rate — a 1-D tensor or sequence of length reservoir_size — assigns each unit its own timescale (heterogeneous reservoir, matching RC.jl's leak_coefficient vector). Every entry must be finite and in (0, 1]; a wrong length or any out-of-range / non-finite entry raises ValueError. A vector leak is stored as a registered buffer (_leak_rate_vec), so it moves with .to(device) / .double() and round-trips through save/load; a scalar stays a plain float with byte-identical dynamics to the pre-vector cell. TYPE: float or 1-D tensor/sequence of shape ``(reservoir_size,)`` DEFAULT: 1.0
noise	Standard deviation of additive Gaussian state noise injected after the activation, following the classical ESN regularizer (Jaeger 2001; Lukoševičius 2012): .. math:: \tilde{h}_t = f(\cdot) + \nu\,\epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, 1) where :math:\nu = noise. Noise is applied only in training mode (self.training is True, like dropout) and is a no-op under :meth:~torch.nn.Module.eval. The default 0.0 disables it entirely, leaving outputs bit-identical to the noiseless cell. The noise stream is reproducible: it is seeded deterministically from seed (independently of the weight-init draws), so a given seed reproduces the same perturbations on every run. Must be non-negative. TYPE: float DEFAULT: 0.0
trainable	If True, reservoir weights are trainable via backpropagation. Standard ESNs use frozen (non-trainable) weights. TYPE: bool DEFAULT: False
feedback_initializer	Initializer for the feedback weight matrix. Accepts a registry name, (name, params), any matrix-building callable, a (callable, params) tuple, or a configured initializer object. TYPE: str, callable, tuple, or InputFeedbackInitializer DEFAULT: None
input_initializer	Initializer for the input weight matrix. Same formats as feedback_initializer. Only used when input_size is provided. TYPE: str, callable, tuple, or InputFeedbackInitializer DEFAULT: None
topology	Structure of the recurrent weight matrix: a registry name (graph or matrix topology), any matrix-building callable, a (callable, params) tuple, or a configured topology object. TYPE: str, callable, tuple, or TopologyInitializer DEFAULT: None
seed	Reproducibility seed that deterministically fixes every reservoir parameter — the recurrent (topology) matrix, the feedback and input weights (including the default uniform(-1, 1) draw used when no initializer is given), and the random bias. Accepts a plain int or a :class:torch.Generator (an int is extracted from the generator's initial_seed() for the NumPy-backed topology/named-initializer path, while the generator itself drives the torch default-init draws). With seed set, a string- or callable-form topology='erdos_renyi' produces the same recurrent matrix on every build, without the ('erdos_renyi', {'seed': ...}) tuple form. An explicit seed inside a tuple/object spec always wins over this argument. When seed=None (the default), string-form graph topologies and the default torch draws are still reproducible under torch.manual_seed because their generators are derived from torch's global RNG. TYPE: int or Generator DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight_feedback	Feedback weight matrix of shape (reservoir_size, feedback_size). TYPE: Parameter
weight_input	Input weight matrix of shape (reservoir_size, input_size), or None if input_size was not provided. TYPE: Parameter or None
weight_hh	Recurrent weight matrix of shape (reservoir_size, reservoir_size). TYPE: Parameter
bias_h	Bias vector of shape (reservoir_size,), or None if bias=False. TYPE: Parameter or None

See Also

resdag.layers.esn.ESNLayer : Layer that sequences this cell. resdag.layers.base.ReservoirCell : Abstract cell interface.

Source code in src/resdag/layers/cells/esn_cell.py

def __init__(
    self,
    reservoir_size: int,
    feedback_size: int,
    input_size: int | None = None,
    spectral_radius: float | None = None,
    bias: bool = True,
    bias_scaling: float = 1.0,
    activation: str = "tanh",
    leak_rate: float | torch.Tensor | Sequence[float] = 1.0,
    noise: float = 0.0,
    trainable: bool = False,
    feedback_initializer: InitializerSpec = None,
    input_initializer: InitializerSpec = None,
    topology: TopologySpec = None,
    seed: SeedLike = None,
) -> None:
    super().__init__()

    # leak_rate / noise are validated by their property setters (below), so
    # both the constructor path (``self.leak_rate = ...``) and live writes
    # (``layer.leak_rate = ...``) share one validation.

    # Store configuration
    self.reservoir_size = reservoir_size
    self.feedback_size = feedback_size
    # Treat input_size == 0 the same as input_size is None (no driving
    # input weight matrix is created).  This avoids the historical
    # (reservoir_size, 0) zero-column tensor produced by passing ``0``
    # explicitly from the premade factories.
    self.input_size = input_size if input_size else None
    self.topology = topology
    self.spectral_radius = spectral_radius
    self.feedback_initializer = feedback_initializer
    self.input_initializer = input_initializer
    # Backing fields for the leak-rate property.  A scalar leak lives in the
    # float ``_leak_rate``; a per-neuron leak lives in the *persistent*
    # buffer ``_leak_rate_vec`` (so it moves with ``.to()`` / ``.double()``
    # and round-trips through ``state_dict()`` / save-load) and flips
    # ``_is_vector_leak`` so the ``!= 1.0`` scalar fast path is bypassed.
    # Both are pre-seeded here so the property setter — shared by the
    # constructor and live writes — can branch on ``_is_vector_leak`` and
    # reassign the buffer freely (buffers accept reassignment, incl. back to
    # ``None``).
    self._leak_rate: float = 1.0
    self._is_vector_leak: bool = False
    self.register_buffer("_leak_rate_vec", None, persistent=True)
    # ``register_buffer(..., None)`` leaves mypy inferring ``None``; annotate
    # the true buffer type so the vector-leak paths type-check.
    self._leak_rate_vec: torch.Tensor | None
    self.leak_rate = leak_rate
    self.noise = noise
    # ``trainable`` is a live property (see below) whose setter flips
    # ``requires_grad_`` on every parameter.  During ``__init__`` the
    # parameters do not exist yet, so stash the requested value in the
    # backing field directly; the actual ``requires_grad`` freeze is applied
    # after ``_initialize_weights`` (see the block at the end of __init__).
    self._trainable = trainable
    self.seed = seed
    # Integer form threaded into the NumPy-backed topology builders and the
    # named feedback/input initializers (they take an int/None seed).  A
    # torch.Generator is reduced to its initial_seed() so the topology stays
    # a pure function of the generator.
    self._seed_int = coerce_seed_to_int(seed)
    # Per-device cache of torch Generators driving the train-mode state
    # noise.  Built lazily on first use (and rebuilt after unpickling — see
    # ``__getstate__``) so the noise stream follows the state's device while
    # staying a deterministic function of ``seed``.  Drawn from a stream
    # *independent* of the weight-init generator so that toggling ``noise``
    # never perturbs weight reproducibility.
    self._noise_generators: dict[torch.device, torch.Generator] = {}

    # Activation function
    self._activation_name = activation
    self.activation_fn = self._get_activation(activation)

    # Store bias config before initialization
    self._bias = bias
    self.bias_scaling = bias_scaling

    # Initialize weight matrices
    self._initialize_weights()

    # Apply the requested trainable state to the freshly-built parameters.
    # Going through the property setter (rather than an inline freeze loop)
    # keeps construction and post-init ``cell.trainable = ...`` writes on a
    # single code path: both flip ``requires_grad_`` on *all* parameters.
    self.trainable = trainable

leak_rate property writable

leak_rate: float | Tensor

float or torch.Tensor : Leaky-integration rate(s) read fresh every step.

A live knob: assigning to it (directly or through ESNLayer.leak_rate = ...) validates the value and takes effect on the next forward. Returns the scalar float for a homogeneous leak (1.0 means no leaking, standard ESN) or the (reservoir_size,) buffer for a per-neuron heterogeneous leak. Every entry stays in (0, 1].

noise property writable

noise: float

float : Std-dev of additive train-mode state noise (>= 0).

A live knob validated on assignment; 0.0 disables noise. Noise is applied only when the cell is in training mode (see _apply_noise).

trainable property writable

trainable: bool

bool : Whether the reservoir weights participate in backpropagation.

This is a live switch, not a construction-time-only flag. Assigning to it flips requires_grad_ on every cell parameter (weight_hh, weight_feedback, weight_input, bias_h) so a post-construction cell.trainable = True genuinely unfreezes the reservoir for a BPTT rollout — the documented reservoir.trainable = True recipe now takes effect instead of being a silent no-op. The getter reflects the current state.

state_size property

state_size: int

Dimensionality of the hidden state vector.

output_size property

output_size: int

Dimensionality of the per-step output (equals state_size for ESN).

activation property

activation: str

str : Name of the activation function.

spectral_radius_achieved property

spectral_radius_achieved: float

float : Realized largest absolute eigenvalue of weight_hh.

Returns the spectral radius actually present in the recurrent matrix after initialization/scaling, as opposed to the requested :attr:spectral_radius target. Computed lazily via the shared :func:resdag.init.topology.estimate_spectral_radius (power iteration / sparse eigs / tiny-N dense fallback), so it stays cheap even for large reservoirs and is GPU-resident for dense matrices.

For a freshly built cell with a non-None spectral_radius this sits within the estimator's tolerance of the target; it differs once the recurrent weights are trained or otherwise modified.

init_state

init_state(batch_size: int, device: device, dtype: dtype) -> Tensor

Return a zero hidden state of shape (batch_size, reservoir_size).

Source code in src/resdag/layers/cells/esn_cell.py

def init_state(
    self,
    batch_size: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Return a zero hidden state of shape ``(batch_size, reservoir_size)``."""
    return torch.zeros(batch_size, self.state_size, device=device, dtype=dtype)

forward

forward(inputs: list[Tensor], state: Tensor) -> tuple[Tensor, Tensor]

Compute the next ESN state for a single timestep.

PARAMETER	DESCRIPTION
inputs	Per-timestep input slices. inputs[0] is the feedback slice of shape (batch, feedback_size). If a driving input is present, inputs[1] has shape (batch, input_size). TYPE: list[Tensor]
state	Current hidden state of shape (batch, reservoir_size). TYPE: Tensor

RETURNS	DESCRIPTION
output	Next hidden state of shape (batch, reservoir_size). TYPE: Tensor
new_state	Same tensor as output (state and output are identical for ESN). TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If the feedback feature dimension does not match self.feedback_size, if a driving input is provided but self.weight_input is None, or if the driving input feature dimension does not match self.input_size.

Notes

When self.noise > 0 and the cell is in training mode, additive Gaussian noise is injected into the post-activation state (see the noise constructor parameter). This matches the noise applied in the :meth:step fast path, so the two paths stay consistent.

Source code in src/resdag/layers/cells/esn_cell.py

def forward(
    self,
    inputs: list[torch.Tensor],
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute the next ESN state for a single timestep.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Per-timestep input slices.  ``inputs[0]`` is the feedback slice
        of shape ``(batch, feedback_size)``.  If a driving input is
        present, ``inputs[1]`` has shape ``(batch, input_size)``.
    state : torch.Tensor
        Current hidden state of shape ``(batch, reservoir_size)``.

    Returns
    -------
    output : torch.Tensor
        Next hidden state of shape ``(batch, reservoir_size)``.
    new_state : torch.Tensor
        Same tensor as output (state and output are identical for ESN).

    Raises
    ------
    ValueError
        If the feedback feature dimension does not match
        ``self.feedback_size``, if a driving input is provided but
        ``self.weight_input`` is ``None``, or if the driving input
        feature dimension does not match ``self.input_size``.

    Notes
    -----
    When ``self.noise > 0`` and the cell is in training mode, additive
    Gaussian noise is injected into the post-activation state (see the
    ``noise`` constructor parameter).  This matches the noise applied in
    the :meth:`step` fast path, so the two paths stay consistent.
    """
    projected = self.project_inputs(inputs)
    recurrent_contrib = F.linear(state, self.weight_hh)
    new_state = self.activation_fn(projected + recurrent_contrib)
    new_state = self._apply_noise(new_state)
    new_state = self._apply_leak(state, new_state)
    return new_state, new_state

project_inputs

project_inputs(inputs: list[Tensor]) -> Tensor

Precompute all input-dependent pre-activation terms.

Computes W_fb x_fb + W_in x_in + b for every timestep at once. Works on full (batch, timesteps, features) sequences (the layer's fast path) and on single-step (batch, features) slices (reused by :meth:forward).

PARAMETER	DESCRIPTION
inputs	Input streams; inputs[0] is feedback, inputs[1] (if present) is the driving input. TYPE: list[Tensor]

RETURNS	DESCRIPTION
Tensor	Input projection with the same leading dimensions as the inputs and trailing dimension reservoir_size.

RAISES	DESCRIPTION
ValueError	If feature dimensions do not match the cell configuration, or a driving input is supplied to a cell built without input_size.

Source code in src/resdag/layers/cells/esn_cell.py

def project_inputs(self, inputs: list[torch.Tensor]) -> torch.Tensor:
    """
    Precompute all input-dependent pre-activation terms.

    Computes ``W_fb x_fb + W_in x_in + b`` for every timestep at once.
    Works on full ``(batch, timesteps, features)`` sequences (the layer's
    fast path) and on single-step ``(batch, features)`` slices (reused by
    :meth:`forward`).

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Input streams; ``inputs[0]`` is feedback, ``inputs[1]`` (if
        present) is the driving input.

    Returns
    -------
    torch.Tensor
        Input projection with the same leading dimensions as the inputs
        and trailing dimension ``reservoir_size``.

    Raises
    ------
    ValueError
        If feature dimensions do not match the cell configuration, or a
        driving input is supplied to a cell built without ``input_size``.
    """
    fb = inputs[0]

    if fb.shape[-1] != self.feedback_size:
        raise ValueError(
            f"Feedback size mismatch. Expected {self.feedback_size}, got {fb.shape[-1]}"
        )

    projected = F.linear(fb, self.weight_feedback)

    if len(inputs) > 1:
        if self.weight_input is None:
            raise ValueError(
                "Reservoir was initialized without input_size, "
                "but driving input was provided in forward pass"
            )
        x = inputs[1]
        if x.shape[-1] != self.input_size:
            raise ValueError(
                f"Driving input size mismatch. Expected {self.input_size}, got {x.shape[-1]}"
            )
        projected = projected + F.linear(x, self.weight_input)

    if self.bias_h is not None:
        projected = projected + self.bias_h

    return projected

step

step(projected_t: Tensor, state: Tensor) -> tuple[Tensor, Tensor]

Recurrent-only single-step update for the projected fast path.

addmm fuses the recurrent matmul with the precomputed input projection into one kernel; with activation and leak that is three kernel launches per timestep instead of six.

PARAMETER	DESCRIPTION
projected_t	Slice of :meth:project_inputs output, shape (batch, reservoir_size). TYPE: Tensor
state	Current hidden state, shape (batch, reservoir_size). TYPE: Tensor

RETURNS	DESCRIPTION
output	Next hidden state. TYPE: Tensor
new_state	Same tensor as output. TYPE: Tensor

Source code in src/resdag/layers/cells/esn_cell.py

def step(
    self,
    projected_t: torch.Tensor,
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Recurrent-only single-step update for the projected fast path.

    ``addmm`` fuses the recurrent matmul with the precomputed input
    projection into one kernel; with activation and leak that is three
    kernel launches per timestep instead of six.

    Parameters
    ----------
    projected_t : torch.Tensor
        Slice of :meth:`project_inputs` output, shape
        ``(batch, reservoir_size)``.
    state : torch.Tensor
        Current hidden state, shape ``(batch, reservoir_size)``.

    Returns
    -------
    output : torch.Tensor
        Next hidden state.
    new_state : torch.Tensor
        Same tensor as output.
    """
    pre_activation = torch.addmm(projected_t, state, self.weight_hh.t())
    new_state = self.activation_fn(pre_activation)
    new_state = self._apply_noise(new_state)
    new_state = self._apply_leak(state, new_state)
    return new_state, new_state

NGCell

NGCell(input_dim: int, k: int = 2, s: int = 1, p: int = 2, cumulative: bool = False, include_constant: bool = True, include_linear: bool = True)

Bases: ReservoirCell

Single-timestep NG-RC feature construction cell.

This is NOT a recurrent cell — it owns no weight matrices and has no recurrent dynamics. The "state" is a FIFO delay buffer of the last (k-1)*s input vectors. The cell wraps feedforward feature construction (time-delayed inputs + polynomial monomials) behind a forward(x, state) -> (features, new_state) interface so that it composes with the same DAG infrastructure as :class:ESNCell.

PARAMETER	DESCRIPTION
input_dim	Dimensionality of a single input vector. TYPE: int
k	Number of delay taps (including the current input). k=1 means only the current input is used (no delay buffer). TYPE: int DEFAULT: 2
s	Spacing between delay taps, in timesteps. TYPE: int DEFAULT: 1
p	Polynomial degree for nonlinear feature construction. See cumulative for the two supported degree conventions. Exact-degree convention (cumulative=False, the default): the nonlinear block contains all unique monomials of exactly degree p from the D = input_dim * k linear features — there are C(D + p - 1, p) of them. Lower-order cross terms (degrees 2, ..., p-1) are not included. For example, NGCell(p=3, include_linear=True) emits constant + linear + cubic with no quadratic terms. This is the historical resdag behaviour and matches Gauthier et al. (arXiv:2106.07688v2), whose Lorenz63 (Eq. 9) and double-scroll (Eq. 10) bases each use a single polynomial degree. Cumulative-degree convention (cumulative=True): the nonlinear block contains every monomial of degree 2, ..., p (or 1, ..., p when include_linear=False), i.e. the full polynomial basis up to degree p. Many NVAR / NG-RC implementations (and configs ported from them) expect this "degree up to p" basis. .. note:: When p == 1 the degree-1 monomials are the linear features O_lin. To avoid emitting two identical blocks (which makes the readout design matrix rank-deficient), the nonlinear block is omitted whenever p == 1 and include_linear is True. With p == 1 and include_linear=False the degree-1 monomials are kept, since they are then the only delay-embedded features. More generally, whenever include_linear is True the degree-1 monomials are dropped from the nonlinear block (in both conventions), since they would duplicate O_lin; the cumulative basis therefore spans degrees 2, ..., p in that case. TYPE: int DEFAULT: 2
cumulative	Degree convention for the nonlinear block (see p). False keeps only exact-degree-p monomials (default, unchanged behaviour); True includes all monomials of every degree up to p (the cumulative / "degree up to p" basis). TYPE: bool DEFAULT: False
include_constant	Whether to prepend a constant 1.0 feature to the output vector. Set True for Lorenz63 forecasting (Eq. 9). TYPE: bool DEFAULT: True
include_linear	Whether to include the linear delay-embedded features O_lin in the output. Set True for both Lorenz63 and double-scroll (Eq. 9, 10). TYPE: bool DEFAULT: True

ATTRIBUTE	DESCRIPTION
feature_dim	Total dimension of the output feature vector O_total:: D = input_dim * k # Degrees emitted in the nonlinear block. Degree 1 is excluded # whenever include_linear is True (those monomials == O_lin): d_lo = 2 if include_linear else 1 degrees = range(d_lo, p + 1) if cumulative else [p] degrees = [g for g in degrees if not (g == 1 and include_linear)] n_nonlin = sum(C(D + g - 1, g) for g in degrees) feature_dim = ( int(include_constant) + int(include_linear) * D + n_nonlin ) In the default exact-degree mode (cumulative=False) this reduces to int(include_constant) + int(include_linear) * D + C(D + p - 1, p), with the C(...) term dropped only in the p == 1 / include_linear=True duplicate-column case. TYPE: int
state_size	Number of rows in the delay buffer: (k - 1) * s. When k=1 this is 0 and the buffer is empty. TYPE: int
monomial_indices	Long tensor of shape (n_monomials, p) containing the column indices into O_lin for each monomial. Registered as a buffer. Each row lists the O_lin columns multiplied together to form one monomial. In exact-degree mode every row has p distinct index slots; in cumulative mode lower-degree monomials are right-padded with -1 sentinels (gathered from a prepended 1.0 column at runtime) so all rows share the same width p. TYPE: Tensor
delay_indices	Long tensor of shape (k-1,) containing the row indices into the delay buffer for extracting the delay taps. Registered as a buffer. TYPE: Tensor

Examples:

Basic usage (k=2, s=1, p=2):

>>> cell = NGCell(input_dim=3)
>>> x = torch.randn(4, 3)                          # (batch=4, d=3)
>>> state = cell.init_state(4, x.device, x.dtype)  # (4, 1, 3)
>>> features, new_state = cell([x], state)
>>> features.shape
torch.Size([4, 28])
>>> new_state.shape
torch.Size([4, 1, 3])

No-delay mode (k=1):

>>> cell = NGCell(input_dim=3, k=1)
>>> state = cell.init_state(1, 'cpu', torch.float32)
>>> features, new_state = cell([torch.randn(1, 3)], state)
>>> features.shape
torch.Size([1, 10])

See Also

resdag.layers.reservoirs.ngrc.NGReservoirLayer : Sequence wrapper for this cell.

Source code in src/resdag/layers/cells/ngrc_cell.py

def __init__(
    self,
    input_dim: int,
    k: int = 2,
    s: int = 1,
    p: int = 2,
    cumulative: bool = False,
    include_constant: bool = True,
    include_linear: bool = True,
) -> None:
    super().__init__()

    if k < 1:
        raise ValueError(f"k must be >= 1, got {k}")
    if s < 1:
        raise ValueError(f"s must be >= 1, got {s}")
    if p < 1:
        raise ValueError(f"p must be >= 1, got {p}")

    self.input_dim = input_dim
    self.k = k
    self.s = s
    self.p = p
    self.cumulative = cumulative
    self.include_constant = include_constant
    self.include_linear = include_linear

    # D = total linear feature dimension
    linear_feature_dim = input_dim * k

    # Precompute delay-tap row indices into the buffer (shape: (k-1,)).
    # For j = 1, ..., k-1:  X_{i - j*s} lives at row (k-1-j)*s of the
    # buffer (which holds inputs chronologically oldest-first).
    if k > 1:
        delay_idx = torch.tensor([(k - 1 - j) * s for j in range(1, k)], dtype=torch.long)
    else:
        delay_idx = torch.zeros(0, dtype=torch.long)
    self.register_buffer("delay_indices", delay_idx)

    # ------------------------------------------------------------------
    # Decide which monomial degrees the nonlinear block emits.
    #
    # Exact-degree mode (cumulative=False): only degree p.
    # Cumulative mode (cumulative=True):    every degree 1..p.
    #
    # When include_linear is True the degree-1 monomials are exactly the
    # columns of O_lin, so emitting them again would duplicate every linear
    # column and make the readout design matrix rank-deficient.  Drop
    # degree 1 from the nonlinear block in that case (this generalises the
    # historical p==1 special case to both conventions).
    # ------------------------------------------------------------------
    degrees = range(1, p + 1) if cumulative else range(p, p + 1)
    emit_degrees = [g for g in degrees if not (g == 1 and include_linear)]

    # Build the monomial index rows.  Exact-degree mode keeps the historical
    # tight (n_monomials, p) layout with no padding.  Cumulative mode mixes
    # degrees, so every row is right-padded to width p with the -1 sentinel,
    # which is gathered from a prepended 1.0 column (the multiplicative
    # identity) at runtime — keeping a single vectorised gather/prod.
    raw_indices: list[tuple[int, ...]] = []
    for g in emit_degrees:
        raw_indices.extend(
            itertools.combinations_with_replacement(range(linear_feature_dim), g)
        )

    self._pad_with_ones: bool = cumulative and len(emit_degrees) > 1
    if self._pad_with_ones:
        # Right-pad short monomials with the -1 ones-sentinel to width p.
        padded = [list(idx) + [-1] * (p - len(idx)) for idx in raw_indices]
        monomial_idx = (
            torch.tensor(padded, dtype=torch.long)
            if padded
            else torch.zeros(0, p, dtype=torch.long)
        )
    else:
        # Tight layout: every monomial already has exactly ``width`` factors.
        width = emit_degrees[0] if emit_degrees else p
        monomial_idx = (
            torch.tensor([list(idx) for idx in raw_indices], dtype=torch.long)
            if raw_indices
            else torch.zeros(0, width, dtype=torch.long)
        )
    self.register_buffer("monomial_indices", monomial_idx)

    n_monomials = len(raw_indices)

    # The nonlinear block is omitted entirely only when no degrees survive
    # the degree-1 filter (i.e. p == 1 and include_linear is True).
    self._emit_nonlinear: bool = len(emit_degrees) > 0

    # Total output dimension
    self._feature_dim: int = (
        int(include_constant)
        + int(include_linear) * linear_feature_dim
        + int(self._emit_nonlinear) * n_monomials
    )

    if self._feature_dim > 10_000:
        warnings.warn(
            f"NGCell: feature_dim={self._feature_dim} exceeds 10,000. "
            f"Consider reducing k, p, or input_dim to avoid combinatorial explosion. "
            f"(linear_feature_dim={linear_feature_dim}, n_monomials={n_monomials})",
            stacklevel=2,
        )

feature_dim property

feature_dim: int

Total dimension of the O_total output vector.

output_size property

output_size: int

Dimensionality of the per-step output (equals feature_dim).

state_size property

state_size: int

Number of rows in the delay buffer: (k-1)*s.

init_state

init_state(batch_size: int, device: device | str, dtype: dtype) -> Tensor

Return a zero-filled initial delay buffer.

PARAMETER	DESCRIPTION
batch_size	Number of samples in the batch. TYPE: int
device	Target device for the buffer. TYPE: device or str
dtype	Target dtype for the buffer. TYPE: dtype

RETURNS	DESCRIPTION
Tensor	Zero tensor of shape (batch_size, state_size, input_dim). When k=1 the second dimension is 0 (empty buffer).

Source code in src/resdag/layers/cells/ngrc_cell.py

def init_state(
    self,
    batch_size: int,
    device: torch.device | str,
    dtype: torch.dtype,
) -> torch.Tensor:
    """
    Return a zero-filled initial delay buffer.

    Parameters
    ----------
    batch_size : int
        Number of samples in the batch.
    device : torch.device or str
        Target device for the buffer.
    dtype : torch.dtype
        Target dtype for the buffer.

    Returns
    -------
    torch.Tensor
        Zero tensor of shape ``(batch_size, state_size, input_dim)``.
        When ``k=1`` the second dimension is 0 (empty buffer).
    """
    return torch.zeros(batch_size, self.state_size, self.input_dim, device=device, dtype=dtype)

forward

forward(inputs: list[Tensor], state: Tensor) -> tuple[Tensor, Tensor]

Compute the NG-RC feature vector for a single timestep.

PARAMETER	DESCRIPTION
inputs	Per-timestep input slices. inputs[0] is the current input of shape (batch, input_dim). NG-RC has no driving inputs, so only inputs[0] is consumed (the wrapping :class:~resdag.layers.reservoirs.ngrc.NGReservoirLayer rejects driving inputs before they reach the cell). TYPE: list[Tensor]
state	Delay buffer of shape (batch, state_size, input_dim). TYPE: Tensor

RETURNS	DESCRIPTION
features	Output feature vector O_total of shape (batch, feature_dim). TYPE: Tensor
new_state	Updated delay buffer of shape (batch, state_size, input_dim). TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If inputs[0]'s trailing (feature) dimension does not equal :attr:input_dim.

Source code in src/resdag/layers/cells/ngrc_cell.py

def forward(
    self,
    inputs: list[torch.Tensor],
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute the NG-RC feature vector for a single timestep.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Per-timestep input slices.  ``inputs[0]`` is the current input
        of shape ``(batch, input_dim)``.  NG-RC has no driving inputs, so
        only ``inputs[0]`` is consumed (the wrapping
        :class:`~resdag.layers.reservoirs.ngrc.NGReservoirLayer` rejects driving
        inputs before they reach the cell).
    state : torch.Tensor
        Delay buffer of shape ``(batch, state_size, input_dim)``.

    Returns
    -------
    features : torch.Tensor
        Output feature vector ``O_total`` of shape
        ``(batch, feature_dim)``.
    new_state : torch.Tensor
        Updated delay buffer of shape ``(batch, state_size, input_dim)``.

    Raises
    ------
    ValueError
        If ``inputs[0]``'s trailing (feature) dimension does not equal
        :attr:`input_dim`.
    """
    x = inputs[0]
    self._check_input_dim(x)
    batch = x.shape[0]

    # ------------------------------------------------------------------
    # 1. Build O_lin (linear delay-embedded features, Eq. 5)
    # ------------------------------------------------------------------
    if self.k > 1:
        # Extract (k-1) delay taps from the buffer *before* update.
        # delay_indices selects rows so that taps[:,0,:] = X_{i-s},
        # taps[:,1,:] = X_{i-2s}, ..., taps[:,k-2,:] = X_{i-(k-1)s}.
        delay_indices = cast(torch.Tensor, self.delay_indices)
        taps = state[:, delay_indices, :]  # (batch, k-1, input_dim)
        # O_lin = [X_i || X_{i-s} || ... || X_{i-(k-1)s}]
        o_lin = torch.cat([x.unsqueeze(1), taps], dim=1).reshape(batch, self.k * self.input_dim)
    else:
        # k=1: no delay taps, just current input
        o_lin = x  # (batch, input_dim)

    # ------------------------------------------------------------------
    # 2. Update delay buffer (FIFO: drop oldest row, append current x)
    # ------------------------------------------------------------------
    if self.state_size > 0:
        # Keep the last state_size rows after appending x.
        new_state = torch.cat([state, x.unsqueeze(1)], dim=1)[:, -self.state_size :, :]
    else:
        new_state = state  # empty tensor, no change

    # ------------------------------------------------------------------
    # 3. Nonlinear features (Eq. 6): polynomial monomials from O_lin.
    #
    # Exact-degree mode emits only degree-p monomials; cumulative mode emits
    # every degree up to p.  The gather/product is delegated to
    # :meth:`_build_nonlinear`, which handles the padded cumulative layout.
    # The block is omitted only when p == 1 and include_linear is True
    # (the degree-1 monomials are then exactly the columns of O_lin).
    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    # 4. Assemble O_total = [c] ⊕ O_lin ⊕ O_nonlin (Eq. 9 / Eq. 10)
    # ------------------------------------------------------------------
    parts: list[torch.Tensor] = []
    if self.include_constant:
        parts.append(torch.ones(batch, 1, device=x.device, dtype=x.dtype))
    if self.include_linear:
        parts.append(o_lin)
    if self._emit_nonlinear:
        parts.append(self._build_nonlinear(o_lin))  # (batch, n_monomials)

    o_total = torch.cat(parts, dim=-1)  # (batch, feature_dim)

    return o_total, new_state

forward_sequence

forward_sequence(x: Tensor, state: Tensor) -> tuple[Tensor, Tensor]

Compute NG-RC features for a whole sequence in a single vectorized pass.

This is the batch-forward fast path used by :meth:~resdag.layers.reservoirs.ngrc.NGReservoirLayer.forward. It is numerically identical to scanning :meth:forward step-by-step (0.0 max-abs-diff, including the warmup region), but contains no per-timestep Python loop: the delay-embedded linear features are built from shifted, front-zero-padded slices of the input, every monomial is gathered over the full (batch, timesteps, *) tensor at once, and the constant-ones block is allocated a single time.

PARAMETER	DESCRIPTION
x	Input sequence of shape (batch, timesteps, input_dim). TYPE: Tensor
state	Incoming delay buffer of shape (batch, state_size, input_dim), chronologically oldest-first. This holds the (k-1)*s inputs immediately preceding x[:, 0] and supplies the delay taps for the first timesteps (zeros for a cold start). TYPE: Tensor

RETURNS	DESCRIPTION
features	Output features O_total of shape (batch, timesteps, feature_dim). TYPE: Tensor
new_state	Updated delay buffer of shape (batch, state_size, input_dim), holding the last state_size inputs of x (or, when the sequence is shorter than the buffer, carried-over history), so a subsequent streaming :meth:forward continues seamlessly. TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If x's trailing (feature) dimension does not equal :attr:input_dim.

See Also

NGCell.forward : Single-step counterpart driving the streaming path.

Source code in src/resdag/layers/cells/ngrc_cell.py

def forward_sequence(
    self,
    x: torch.Tensor,
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute NG-RC features for a whole sequence in a single vectorized pass.

    This is the batch-forward fast path used by
    :meth:`~resdag.layers.reservoirs.ngrc.NGReservoirLayer.forward`.  It is
    numerically identical to scanning :meth:`forward` step-by-step
    (``0.0`` max-abs-diff, including the warmup region), but contains no
    per-timestep Python loop: the delay-embedded linear features are built
    from shifted, front-zero-padded slices of the input, every monomial is
    gathered over the full ``(batch, timesteps, *)`` tensor at once, and the
    constant-ones block is allocated a single time.

    Parameters
    ----------
    x : torch.Tensor
        Input sequence of shape ``(batch, timesteps, input_dim)``.
    state : torch.Tensor
        Incoming delay buffer of shape ``(batch, state_size, input_dim)``,
        chronologically oldest-first.  This holds the ``(k-1)*s`` inputs
        immediately preceding ``x[:, 0]`` and supplies the delay taps for
        the first timesteps (zeros for a cold start).

    Returns
    -------
    features : torch.Tensor
        Output features ``O_total`` of shape
        ``(batch, timesteps, feature_dim)``.
    new_state : torch.Tensor
        Updated delay buffer of shape ``(batch, state_size, input_dim)``,
        holding the last ``state_size`` inputs of ``x`` (or, when the
        sequence is shorter than the buffer, carried-over history), so a
        subsequent streaming :meth:`forward` continues seamlessly.

    Raises
    ------
    ValueError
        If ``x``'s trailing (feature) dimension does not equal
        :attr:`input_dim`.

    See Also
    --------
    NGCell.forward : Single-step counterpart driving the streaming path.
    """
    self._check_input_dim(x)
    batch, seq_len, _ = x.shape
    pad = self.state_size  # == (k - 1) * s

    # ------------------------------------------------------------------
    # 1. Build O_lin for the whole sequence (Eq. 5).
    #
    # Prepend the incoming delay buffer (oldest-first) so the first
    # timesteps draw their delay taps from carried-over history; a cold
    # start contributes zeros, matching the per-step loop's empty buffer.
    # For delay tap j (0..k-1), X_{i-j*s} at output index i lives at
    # padded[:, pad + i - j*s, :], i.e. the length-T window starting at
    # column (pad - j*s).
    # ------------------------------------------------------------------
    if self.k > 1:
        padded = torch.cat([state, x], dim=1)  # (batch, pad + T, input_dim)
        taps = [
            padded[:, pad - j * self.s : pad - j * self.s + seq_len, :] for j in range(self.k)
        ]
        o_lin = torch.cat(taps, dim=-1)  # (batch, T, k * input_dim)
    else:
        # k=1: no delay taps, O_lin is just the current input.
        o_lin = x  # (batch, T, input_dim)

    # ------------------------------------------------------------------
    # 2. Updated delay buffer: the last ``pad`` rows of the padded stream
    #    (FIFO drop-oldest / append-current applied to the full sequence).
    # ------------------------------------------------------------------
    if pad > 0:
        new_state = torch.cat([state, x], dim=1)[:, -pad:, :]
    else:
        new_state = state  # empty buffer, unchanged

    # ------------------------------------------------------------------
    # 3. Assemble O_total = [c] + O_lin + O_nonlin (Eq. 9 / Eq. 10).
    #    The constant-ones block is allocated once for the whole sequence;
    #    monomials are gathered over the full (batch, T, *) tensor in one op.
    # ------------------------------------------------------------------
    parts: list[torch.Tensor] = []
    if self.include_constant:
        parts.append(torch.ones(batch, seq_len, 1, device=x.device, dtype=x.dtype))
    if self.include_linear:
        parts.append(o_lin)
    if self._emit_nonlinear:
        parts.append(self._build_nonlinear(o_lin))  # (batch, T, n_monomials)

    o_total = torch.cat(parts, dim=-1)  # (batch, T, feature_dim)

    return o_total, new_state

validate_state

validate_state(state: Tensor) -> None

Validate the 3-D delay-buffer layout used by NG-RC.

PARAMETER	DESCRIPTION
state	Candidate state of shape (batch, state_size, input_dim). TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If state is not 3-D or its non-batch dimensions do not match (state_size, input_dim).

Source code in src/resdag/layers/cells/ngrc_cell.py

def validate_state(self, state: torch.Tensor) -> None:
    """
    Validate the 3-D delay-buffer layout used by NG-RC.

    Parameters
    ----------
    state : torch.Tensor
        Candidate state of shape ``(batch, state_size, input_dim)``.

    Raises
    ------
    ValueError
        If ``state`` is not 3-D or its non-batch dimensions do not match
        ``(state_size, input_dim)``.
    """
    expected = (self.state_size, self.input_dim)
    if state.dim() != 3 or state.shape[1:] != torch.Size(expected):
        raise ValueError(
            f"NGCell.validate_state: expected a 3-D delay buffer of shape "
            f"(batch, {self.state_size}, {self.input_dim}); "
            f"got tensor of shape {tuple(state.shape)}."
        )

reservoirs

BaseReservoirLayer

BaseReservoirLayer(cell: ReservoirCell)

Bases: Module, ABC

Abstract base that owns the sequence loop and all state-management methods.

Subclasses create a :class:ReservoirCell and pass it to this constructor. The sequence loop (iterating over timesteps and calling the cell) lives here; the cell handles the per-step computation.

PARAMETER	DESCRIPTION
cell	Concrete cell instance that performs the single-step update. TYPE: ReservoirCell

ATTRIBUTE	DESCRIPTION
cell	The wrapped single-step cell. TYPE: ReservoirCell
state	Current reservoir state, or None if not yet initialized. The layout is cell-defined: 2-D (batch, cell.state_size) for :class:~resdag.layers.cells.esn_cell.ESNCell, 3-D (batch, cell.state_size, input_dim) delay buffer for :class:~resdag.layers.cells.ngrc_cell.NGCell. See :attr:ReservoirCell.state_size <resdag.layers.cells.base_cell.ReservoirCell.state_size> for the leading state dimension. TYPE: Tensor or None
detach_state_between_calls	If True (default), the stored state is detached from the autograd graph at the end of each forward call (truncated BPTT at call boundaries). Gradients still flow through the returned states within a call — this only severs gradient flow across calls, which otherwise raises "backward through the graph a second time" when a model is trained with SGD over consecutive batches without resetting the reservoir. Set to False only if you intentionally backprop through state carried across multiple forward calls (and manage the retained graphs yourself). TYPE: bool

See Also

resdag.layers.esn.ESNLayer : Concrete ESN layer built on this base. resdag.layers.base.ReservoirCell : Abstract cell interface.

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def __init__(self, cell: ReservoirCell) -> None:
    super().__init__()
    self.cell = cell
    # Register the reservoir state as a *non-persistent* buffer so that
    # ``nn.Module._apply`` (driven by ``.to()`` / ``.cuda()`` / ``.double()``)
    # moves a warmed-up state along with the module's parameters, instead of
    # leaving it on the original device/dtype and triggering a silent
    # zero-reinit on the next forward.  ``persistent=False`` keeps it out of
    # ``state_dict()`` so the ``save`` / ``include_states`` split is
    # untouched.  The buffer is initialised to ``None`` and reassigned with a
    # concrete tensor on first use (buffers accept reassignment, including
    # back to ``None`` for lazy re-init).
    self.register_buffer("state", None, persistent=False)
    self.state: torch.Tensor | None
    self.detach_state_between_calls: bool = True
    # Distinguishes a state the user restored on purpose (via ``set_state``)
    # from one that merely accumulated through forward passes or lazy
    # zero-init.  Only the former is protected against a silent batch-size
    # re-init in ``_maybe_init_state`` — the latter is free to auto-resize.
    self._state_user_set: bool = False
    # Per-configuration cache of compiled teacher-forced chunk scans, keyed by
    # ``(id(cell), chunk, dtype, device)``.  Populated lazily by
    # ``forward_stateless(compile=...)`` and dropped on pickle / deepcopy (it
    # holds a ``torch.compile`` closure over the live cell that would
    # otherwise close over stale layers on a copy — mirrors the forecast
    # engine's ``_flat_step_cache`` hygiene in ``ESNModel``).
    self._compiled_scan_cache: dict[
        tuple[int, int, torch.dtype, torch.device], _ChunkScanFn
    ] = {}

forward_stateless

forward_stateless(inputs: list[Tensor], state: Tensor, *, compile: bool = False) -> tuple[Tensor, Tensor]

Run the time loop purely: thread state in and out, touch no self state.

This is the functional core of the reservoir. It takes the initial state as an explicit argument, threads it through the per-timestep loop, and returns the per-step outputs together with the final state. It never reads or writes :attr:state, performs no in-place tensor writes (per-step outputs are collected with :func:torch.stack), and the loop body contains no data-dependent Python branch. Those three properties are what make it torch.compile-scan, :func:torch.func.vmap, and TorchScript/ONNX-export friendly, unlike the stateful :meth:forward wrapper that drives it.

PARAMETER	DESCRIPTION
inputs	Feedback tensor first, then at most one driving-input tensor, each of shape (batch, timesteps, features). TYPE: list[Tensor]
state	Initial reservoir state, shape (batch, ...) as produced by self.cell.init_state. TYPE: Tensor
compile	Opt-in compiled/chunked engine for the teacher-forced sequence loop (warmup, ESNTrainer.fit forward, windowed_forecast re-sync). When True and the cell exposes the projected fast path (its project_inputs is not None), the recurrent-only cell.step update is unrolled DEFAULT_COMPILE_CHUNK steps at a time inside a :func:torch.compile (mode="reduce-overhead", fullgraph=True) chunk driven chunk-by-chunk with cudagraph_mark_step_begin and state clones at boundaries — the teacher-forced analogue of the forecast engine (:mod:resdag.core._flat_inference). On the GPU this sheds the per-step launch/dispatch overhead that dominates the eager loop (~6-7x for a N=500/T=2000 warmup). The default (False) keeps the eager time loop below byte-for-byte unchanged. The compiled path is inference-only: it transparently falls back to the eager loop when any parameter requires grad (a differentiable BPTT rollout must stay eager), when there is no projected fast path (e.g. NG-RC), on torch < 2.10, or on any compilation failure (with a :class:RuntimeWarning), so a correct result is always produced. Because chaotic reservoir dynamics under torch.compile need not be bit-exact with eager, the compiled output matches the eager output to a tight allclose (typically < 1e-5 max abs diff; closer for identity/linear activation) rather than bitwise. TYPE: (bool, keyword - only) DEFAULT: False

RETURNS	DESCRIPTION
outputs	Per-step outputs, shape (batch, timesteps, cell.output_size). TYPE: Tensor
new_state	Final reservoir state after the last timestep. TYPE: Tensor

Notes

The cross-call detach applied by :meth:forward (truncated BPTT) lives in that wrapper, outside this method — gradients flow through the returned new_state here, so callers that backpropagate through state carried across calls are free to do so.

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def forward_stateless(
    self,
    inputs: list[torch.Tensor],
    state: torch.Tensor,
    *,
    compile: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Run the time loop purely: thread ``state`` in and out, touch no ``self`` state.

    This is the functional core of the reservoir.  It takes the initial
    state as an explicit argument, threads it through the per-timestep
    loop, and returns the per-step outputs together with the final state.
    It never reads or writes :attr:`state`, performs no in-place tensor
    writes (per-step outputs are collected with :func:`torch.stack`), and
    the loop body contains no data-dependent Python branch.  Those three
    properties are what make it ``torch.compile``-scan, :func:`torch.func.vmap`,
    and TorchScript/ONNX-export friendly, unlike the stateful
    :meth:`forward` wrapper that drives it.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Feedback tensor first, then at most one driving-input tensor, each
        of shape ``(batch, timesteps, features)``.
    state : torch.Tensor
        Initial reservoir state, shape ``(batch, ...)`` as produced by
        ``self.cell.init_state``.
    compile : bool, keyword-only, default=False
        Opt-in compiled/chunked engine for the teacher-forced sequence loop
        (warmup, ``ESNTrainer.fit`` forward, ``windowed_forecast`` re-sync).
        When ``True`` and the cell exposes the projected fast path (its
        ``project_inputs`` is not ``None``), the recurrent-only ``cell.step``
        update is unrolled ``DEFAULT_COMPILE_CHUNK`` steps at a time inside a
        :func:`torch.compile` (``mode="reduce-overhead"``, ``fullgraph=True``)
        chunk driven chunk-by-chunk with ``cudagraph_mark_step_begin`` and
        state clones at boundaries — the teacher-forced analogue of the
        forecast engine (:mod:`resdag.core._flat_inference`).  On the GPU this
        sheds the per-step launch/dispatch overhead that dominates the eager
        loop (~6-7x for a ``N=500``/``T=2000`` warmup).  The default (``False``)
        keeps the eager time loop below byte-for-byte unchanged.

        The compiled path is **inference-only**: it transparently falls back
        to the eager loop when any parameter requires grad (a differentiable
        BPTT rollout must stay eager), when there is no projected fast path
        (e.g. NG-RC), on ``torch < 2.10``, or on any compilation failure
        (with a :class:`RuntimeWarning`), so a correct result is always
        produced.  Because chaotic reservoir dynamics under ``torch.compile``
        need not be bit-exact with eager, the compiled output matches the
        eager output to a tight ``allclose`` (typically ``< 1e-5`` max abs
        diff; closer for identity/linear activation) rather than bitwise.

    Returns
    -------
    outputs : torch.Tensor
        Per-step outputs, shape ``(batch, timesteps, cell.output_size)``.
    new_state : torch.Tensor
        Final reservoir state after the last timestep.

    Notes
    -----
    The cross-call detach applied by :meth:`forward` (truncated BPTT) lives
    in that wrapper, *outside* this method — gradients flow through the
    returned ``new_state`` here, so callers that backpropagate through state
    carried across calls are free to do so.
    """
    feedback = inputs[0]
    seq_len = feedback.shape[1]

    # Fast path: cells that split their update into input-dependent and
    # state-dependent parts precompute the former for the whole sequence
    # in one batched matmul, leaving only the recurrent term in the loop.
    # Choosing the path is a per-cell static decision made once, outside
    # the loop, so the loop body itself stays branch-free.
    projected = self.cell.project_inputs(inputs)

    # Opt-in compiled/chunked scan.  Gated to the projected fast path and
    # inference-only (no grad required): otherwise the eager loop below runs
    # exactly as before, keeping the default path byte-for-byte unchanged.
    if compile and projected is not None and seq_len > 0 and self._compile_eligible(projected):
        compiled = self._maybe_compiled_scan(projected, state)
        if compiled is not None:
            return self._run_compiled_scan(compiled, projected, state)

    outputs: list[torch.Tensor] = []
    if projected is not None:
        for t in range(seq_len):
            output, state = self.cell.step(projected[:, t, :], state)
            outputs.append(output)
    else:
        for t in range(seq_len):
            inputs_t = [stream[:, t, :] for stream in inputs]
            output, state = self.cell(inputs_t, state)
            outputs.append(output)

    if not outputs:
        # Degenerate empty sequence (T == 0): torch.stack rejects an empty
        # list, so return the (batch, 0, output_size) tensor the previous
        # preallocating loop produced, with the input ``state`` untouched.
        empty = feedback.new_empty(feedback.shape[0], 0, self.cell.output_size)
        return empty, state

    return torch.stack(outputs, dim=1), state

step_stateless

step_stateless(inputs: list[Tensor], state: Tensor) -> tuple[Tensor, Tensor]

Single-timestep analogue of :meth:forward_stateless: state in, state out.

Where :meth:forward_stateless consumes whole (batch, timesteps, features) sequences and runs the time loop, this consumes a single (batch, features) slice per stream and performs exactly one cell update. It is the per-step primitive that the flattened autoregressive engine (:meth:resdag.core.ESNModel.forecast) drives, sidestepping the sequence-loop bookkeeping (torch.stack, shape[1] reads) that :meth:forward / :meth:forward_stateless pay even for a length-1 sequence.

PARAMETER	DESCRIPTION
inputs	Feedback slice first, then at most one driving-input slice, each of shape (batch, features). TYPE: list[Tensor]
state	Current reservoir state, shape (batch, ...) as produced by self.cell.init_state. TYPE: Tensor

RETURNS	DESCRIPTION
output	Per-step output, shape (batch, cell.output_size). TYPE: Tensor
new_state	Updated reservoir state. TYPE: Tensor

Notes

The project-vs-fallback choice is the same per-cell static decision made in :meth:forward_stateless (no data-dependent branch in the hot path), so this method stays torch.compile-fullgraph friendly. Like :meth:forward_stateless it never reads or writes :attr:state, and it applies no cross-call detach — the forecast engine runs under :func:torch.no_grad, so the threaded state never carries a graph.

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def step_stateless(
    self,
    inputs: list[torch.Tensor],
    state: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Single-timestep analogue of :meth:`forward_stateless`: state in, state out.

    Where :meth:`forward_stateless` consumes whole
    ``(batch, timesteps, features)`` sequences and runs the time loop, this
    consumes a single ``(batch, features)`` slice per stream and performs
    exactly one cell update.  It is the per-step primitive that the
    flattened autoregressive engine
    (:meth:`resdag.core.ESNModel.forecast`) drives, sidestepping the
    sequence-loop bookkeeping (``torch.stack``, ``shape[1]`` reads) that
    :meth:`forward` / :meth:`forward_stateless` pay even for a length-1
    sequence.

    Parameters
    ----------
    inputs : list[torch.Tensor]
        Feedback slice first, then at most one driving-input slice, each of
        shape ``(batch, features)``.
    state : torch.Tensor
        Current reservoir state, shape ``(batch, ...)`` as produced by
        ``self.cell.init_state``.

    Returns
    -------
    output : torch.Tensor
        Per-step output, shape ``(batch, cell.output_size)``.
    new_state : torch.Tensor
        Updated reservoir state.

    Notes
    -----
    The project-vs-fallback choice is the same per-cell static decision made
    in :meth:`forward_stateless` (no data-dependent branch in the hot path),
    so this method stays ``torch.compile``-fullgraph friendly.  Like
    :meth:`forward_stateless` it never reads or writes :attr:`state`, and it
    applies no cross-call detach — the forecast engine runs under
    :func:`torch.no_grad`, so the threaded state never carries a graph.
    """
    projected = self.cell.project_inputs(inputs)
    if projected is not None:
        return self.cell.step(projected, state)
    # ``nn.Module.__call__`` is typed ``Any``; unpack then return a tuple
    # literal so the declared ``tuple[Tensor, Tensor]`` is honoured (mirrors
    # ``forward_stateless``).
    output, new_state = self.cell(inputs, state)
    return output, new_state

forward

forward(feedback: Tensor, *driving_inputs: Tensor, compile: bool = False) -> Tensor

Process an input sequence through the reservoir.

Computes reservoir states for each timestep using the feedback signal and optional driving inputs. This is a thin stateful wrapper over the pure :meth:forward_stateless: it validates inputs, lazily initialises :attr:state, runs the functional loop, stores the returned state, and applies the cross-call detach.

PARAMETER	DESCRIPTION
feedback	Feedback signal of shape (batch, timesteps, feedback_size). TYPE: Tensor
*driving_inputs	Optional driving input of shape (batch, timesteps, input_size). At most one driving input tensor is supported. TYPE: Tensor DEFAULT: ()
compile	Route the teacher-forced sequence loop through the opt-in compiled/chunked engine. Forwarded verbatim to :meth:forward_stateless — see its compile parameter for the eligibility gating (projected fast path, inference-only) and the transparent eager fallback. The default (False) is byte-for-byte identical to the historical eager path. TYPE: (bool, keyword - only) DEFAULT: False

RETURNS	DESCRIPTION
Tensor	Reservoir states for all timesteps, shape (batch, timesteps, cell.output_size).

RAISES	DESCRIPTION
ValueError	If feedback is not 3-D, if more than one driving input is supplied, or if the driving input batch/sequence dimensions do not match feedback.

Notes

The layer maintains internal state across forward calls. Use :meth:reset_state to clear the state between independent sequences. For a pure, self-free variant (used by the compile-scan, vmap, and export paths) call :meth:forward_stateless directly.

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def forward(
    self,
    feedback: torch.Tensor,
    *driving_inputs: torch.Tensor,
    compile: bool = False,
) -> torch.Tensor:
    """
    Process an input sequence through the reservoir.

    Computes reservoir states for each timestep using the feedback
    signal and optional driving inputs.  This is a thin stateful wrapper
    over the pure :meth:`forward_stateless`: it validates inputs, lazily
    initialises :attr:`state`, runs the functional loop, stores the
    returned state, and applies the cross-call detach.

    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)``.
        At most one driving input tensor is supported.
    compile : bool, keyword-only, default=False
        Route the teacher-forced sequence loop through the opt-in
        compiled/chunked engine.  Forwarded verbatim to
        :meth:`forward_stateless` — see its ``compile`` parameter for the
        eligibility gating (projected fast path, inference-only) and the
        transparent eager fallback.  The default (``False``) is byte-for-byte
        identical to the historical eager path.

    Returns
    -------
    torch.Tensor
        Reservoir states for all timesteps, shape
        ``(batch, timesteps, cell.output_size)``.

    Raises
    ------
    ValueError
        If ``feedback`` is not 3-D, if more than one driving input is
        supplied, or if the driving input batch/sequence dimensions do not
        match ``feedback``.

    Notes
    -----
    The layer maintains internal state across forward calls.  Use
    :meth:`reset_state` to clear the state between independent sequences.
    For a pure, ``self``-free variant (used by the compile-scan, vmap, and
    export paths) call :meth:`forward_stateless` directly.
    """
    if feedback.dim() != 3:
        raise ValueError(f"Feedback must be 3D (B, T, F), got shape {feedback.shape}")

    batch_size, seq_len, _ = feedback.shape

    if len(driving_inputs) > 0:
        if len(driving_inputs) > 1:
            raise ValueError("Only one driving input tensor allowed")
        driving_input = driving_inputs[0]
        if driving_input.shape[0] != batch_size or driving_input.shape[1] != seq_len:
            raise ValueError(
                f"Driving input must match feedback dimensions. "
                f"Feedback: {feedback.shape}, Driving: {driving_input.shape}"
            )

    self._maybe_init_state(batch_size, feedback.device, feedback.dtype)
    assert self.state is not None  # guaranteed by _maybe_init_state; narrows for mypy

    outputs, new_state = self.forward_stateless(
        [feedback, *driving_inputs], self.state, compile=compile
    )

    # Truncated BPTT at call boundaries: gradients flow through the returned
    # states within this call, but the *stored* state must not keep the
    # graph alive — a later forward + backward would otherwise try to
    # backprop through this call's already-freed graph.  This data-dependent
    # ``grad_fn`` branch lives here, in the wrapper, never inside the pure
    # loop of ``forward_stateless``.
    if self.detach_state_between_calls and new_state.grad_fn is not None:
        new_state = new_state.detach()

    self.state = new_state
    # The stored state has now evolved one or more steps past whatever the
    # user restored; it is no longer the as-set tensor, so drop the
    # user-set protection — a subsequent batch-size change is back to the
    # ordinary auto-resize regime.
    self._state_user_set = False
    return outputs

reference_device_dtype

reference_device_dtype() -> tuple[device, dtype]

Resolve the canonical device and dtype of this reservoir's weights.

The reference is the first floating-point parameter or buffer of the inner cell (the same scan used when lazily allocating a zero state), falling back to cpu / float32 for weightless cells such as the NG-RC reservoir. This is the device/dtype an incoming forward pass is expected to use, so callers restoring a saved state should coerce it to this reference to avoid a silent re-initialisation in :meth:_maybe_init_state.

RETURNS	DESCRIPTION
device	Device of the cell's first floating-point tensor, or cpu. TYPE: device
dtype	Dtype of the cell's first floating-point tensor, or float32. TYPE: dtype

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def reference_device_dtype(self) -> tuple[torch.device, torch.dtype]:
    """
    Resolve the canonical device and dtype of this reservoir's weights.

    The reference is the first floating-point parameter or buffer of the
    inner cell (the same scan used when lazily allocating a zero state),
    falling back to ``cpu`` / ``float32`` for weightless cells such as the
    NG-RC reservoir.  This is the device/dtype an incoming forward pass is
    expected to use, so callers restoring a saved state should coerce it to
    this reference to avoid a silent re-initialisation in
    :meth:`_maybe_init_state`.

    Returns
    -------
    device : torch.device
        Device of the cell's first floating-point tensor, or ``cpu``.
    dtype : torch.dtype
        Dtype of the cell's first floating-point tensor, or ``float32``.
    """
    ref = next(
        (
            t
            for t in chain(self.cell.parameters(), self.cell.buffers())
            if t.is_floating_point()
        ),
        None,
    )
    device = ref.device if ref is not None else torch.device("cpu")
    dtype = ref.dtype if ref is not None else torch.float32
    return device, dtype

reset_state

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

Reset internal state to zero.

PARAMETER	DESCRIPTION
batch_size	If provided, initialize state with this batch size using the cell's device and dtype. If None, state is set to None and will be lazily initialized on the next forward pass. TYPE: int DEFAULT: None

Examples:

>>> layer.reset_state()          # Lazy initialization
>>> layer.reset_state(batch_size=4)  # Explicit zero state

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def reset_state(self, batch_size: int | None = None) -> None:
    """
    Reset internal state to zero.

    Parameters
    ----------
    batch_size : int, optional
        If provided, initialize state with this batch size using the
        cell's device and dtype.  If ``None``, state is set to ``None``
        and will be lazily initialized on the next forward pass.

    Examples
    --------
    >>> layer.reset_state()          # Lazy initialization
    >>> layer.reset_state(batch_size=4)  # Explicit zero state
    """
    if batch_size is not None:
        if self.state is not None:
            device, dtype = self.state.device, self.state.dtype
        else:
            device, dtype = self.reference_device_dtype()
        self.state = self.cell.init_state(batch_size, device, dtype)
    else:
        self.state = None
    # A reset (zeros or back-to-None) is a deliberate opt-in to automatic
    # batch-size re-initialisation: drop any user-set protection.
    self._state_user_set = False

get_state

get_state() -> Tensor | None

Get a copy of the current internal state.

RETURNS	DESCRIPTION
Tensor or None	Clone of the current state tensor, or None if not yet initialized. The layout is cell-defined: 2-D (batch, cell.state_size) for :class:~resdag.layers.cells.esn_cell.ESNCell, 3-D (batch, cell.state_size, input_dim) delay buffer for :class:~resdag.layers.cells.ngrc_cell.NGCell.

Examples:

>>> state = layer.get_state()
>>> if state is not None:
...     print(f"State shape: {state.shape}")

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def get_state(self) -> torch.Tensor | None:
    """
    Get a copy of the current internal state.

    Returns
    -------
    torch.Tensor or None
        Clone of the current state tensor, or ``None`` if not yet
        initialized.  The layout is cell-defined: 2-D
        ``(batch, cell.state_size)`` for
        :class:`~resdag.layers.cells.esn_cell.ESNCell`, 3-D
        ``(batch, cell.state_size, input_dim)`` delay buffer for
        :class:`~resdag.layers.cells.ngrc_cell.NGCell`.

    Examples
    --------
    >>> state = layer.get_state()
    >>> if state is not None:
    ...     print(f"State shape: {state.shape}")
    """
    return self.state.clone() if self.state is not None else None

set_state

set_state(state: Tensor) -> None

Set the internal state to a specific tensor.

Validation is delegated to self.cell.validate_state(state) so each cell type owns its own state-shape contract (2-D for ESNCell, 3-D delay buffer for NGCell, etc.).

Batch-size contract

The restored state's batch size (state.shape[0]) pins the batch size the next forward pass must use. A restored state is treated as deliberate: if a subsequent forward is called with a different batch size, :meth:_maybe_init_state raises a :class:RuntimeError instead of silently zero-reinitialising and discarding the state you restored. To forecast at a different batch size, restore (or :meth:reset_state and re-warm) a state of the matching batch size; to opt back into automatic batch-size re-initialisation, call :meth:reset_state. Device and dtype are not pinned — the state moves with the module under .to() and is cast to match an incoming forward.

Autograd contract

The stored state honours :attr:detach_state_between_calls, exactly as the cross-call truncated-BPTT detach in :meth:forward does. Under the default (True) a restored state is treated as a fresh initial condition: it is stored as state.detach().clone() so it carries no grad_fn and a later forward + backward cannot try to backprop through state's already-freed graph (the "backward through the graph a second time" / retained-memory failure that :attr:detach_state_between_calls exists to prevent). Set detach_state_between_calls=False to store state.clone() with its graph intact, when you intentionally backprop through state carried in across calls (and manage the retained graph yourself).

PARAMETER	DESCRIPTION
state	New state tensor. Its batch size (shape[0]) becomes the required batch size of the next forward pass. TYPE: Tensor

RAISES	DESCRIPTION
ValueError	If the cell's :meth:ReservoirCell.validate_state rejects the tensor. The error includes the cell class name and the offending shape.

See Also

reset_state : Clear the state and opt back into automatic batch-size re-initialisation.

Examples:

>>> saved = layer.get_state()      # captured at batch size B
>>> # ... process data ...
>>> layer.set_state(saved)         # next forward must also use batch B

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def set_state(self, state: torch.Tensor) -> None:
    """
    Set the internal state to a specific tensor.

    Validation is delegated to ``self.cell.validate_state(state)`` so
    each cell type owns its own state-shape contract (2-D for ESNCell,
    3-D delay buffer for NGCell, etc.).

    Batch-size contract
    -------------------
    The restored state's batch size (``state.shape[0]``) pins the batch size
    the **next forward pass must use**.  A restored state is treated as
    deliberate: if a subsequent ``forward`` is called with a different batch
    size, :meth:`_maybe_init_state` raises a :class:`RuntimeError` instead of
    silently zero-reinitialising and discarding the state you restored.  To
    forecast at a different batch size, restore (or :meth:`reset_state` and
    re-warm) a state of the matching batch size; to opt back into automatic
    batch-size re-initialisation, call :meth:`reset_state`.  Device and dtype
    are not pinned — the state moves with the module under ``.to()`` and is
    cast to match an incoming forward.

    Autograd contract
    -----------------
    The stored state honours :attr:`detach_state_between_calls`, exactly as
    the cross-call truncated-BPTT detach in :meth:`forward` does.  Under the
    default (``True``) a restored state is treated as a fresh initial
    condition: it is stored as ``state.detach().clone()`` so it carries no
    ``grad_fn`` and a later forward + backward cannot try to backprop
    through ``state``'s already-freed graph (the "backward through the graph
    a second time" / retained-memory failure that
    :attr:`detach_state_between_calls` exists to prevent).  Set
    ``detach_state_between_calls=False`` to store ``state.clone()`` with its
    graph intact, when you intentionally backprop through state carried in
    across calls (and manage the retained graph yourself).

    Parameters
    ----------
    state : torch.Tensor
        New state tensor.  Its batch size (``shape[0]``) becomes the
        required batch size of the next forward pass.

    Raises
    ------
    ValueError
        If the cell's :meth:`ReservoirCell.validate_state` rejects the
        tensor.  The error includes the cell class name and the
        offending shape.

    See Also
    --------
    reset_state : Clear the state and opt back into automatic batch-size
        re-initialisation.

    Examples
    --------
    >>> saved = layer.get_state()      # captured at batch size B
    >>> # ... process data ...
    >>> layer.set_state(saved)         # next forward must also use batch B
    """
    try:
        self.cell.validate_state(state)
    except ValueError as exc:
        # Re-raise with a richer message including a recovery hint.
        raise ValueError(
            f"{type(self).__name__}.set_state failed: {exc}. "
            f"Tip: call reset_state(batch_size=...) before set_state() "
            f"to materialise an empty state of the expected layout."
        ) from None
    # Honour the truncated-BPTT contract on restore, just as ``forward``
    # does for the evolved state (see lines around the cross-call detach
    # there): a restored state is a fresh initial condition by default, so
    # under ``detach_state_between_calls`` strip any incoming ``grad_fn``
    # before storing it — otherwise restoring a graph-bearing tensor leaks
    # an autograd graph across calls (#146).  ``clone`` always runs so the
    # stored state is an independent copy regardless of the detach choice.
    if self.detach_state_between_calls and state.grad_fn is not None:
        self.state = state.detach().clone()
    else:
        self.state = state.clone()
    # Mark this as a deliberate restore: a later forward at a mismatched
    # batch size must raise rather than silently zero-reinit (see #145).
    self._state_user_set = True

set_random_state

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

Set the internal state to random (standard-normal) values.

PARAMETER	DESCRIPTION
batch_size	If provided, lazily initialise the state with this batch size before filling it with random values. Uses device / dtype kwargs when given, otherwise the cell's first floating-point parameter/buffer (falling back to CPU/float32). TYPE: int DEFAULT: None
device	Target device when lazily initialising. Ignored if the state is already initialised. TYPE: device DEFAULT: None
dtype	Target dtype when lazily initialising. Ignored if the state is already initialised. TYPE: dtype DEFAULT: None

RAISES	DESCRIPTION
RuntimeError	If the state has not been initialised and batch_size is not provided.

Examples:

>>> # Lazy: initialise then randomise in one call
>>> layer.set_random_state(batch_size=4)
>>> # Already-initialised state: just randomise in place
>>> layer.set_random_state()

Source code in src/resdag/layers/reservoirs/base_reservoir.py

def set_random_state(
    self,
    batch_size: int | None = None,
    device: torch.device | None = None,
    dtype: torch.dtype | None = None,
) -> None:
    """
    Set the internal state to random (standard-normal) values.

    Parameters
    ----------
    batch_size : int, optional
        If provided, lazily initialise the state with this batch size
        before filling it with random values.  Uses ``device`` / ``dtype``
        kwargs when given, otherwise the cell's first floating-point
        parameter/buffer (falling back to CPU/float32).
    device : torch.device, optional
        Target device when lazily initialising.  Ignored if the state is
        already initialised.
    dtype : torch.dtype, optional
        Target dtype when lazily initialising.  Ignored if the state is
        already initialised.

    Raises
    ------
    RuntimeError
        If the state has not been initialised and ``batch_size`` is not
        provided.

    Examples
    --------
    >>> # Lazy: initialise then randomise in one call
    >>> layer.set_random_state(batch_size=4)
    >>> # Already-initialised state: just randomise in place
    >>> layer.set_random_state()
    """
    if self.state is None:
        if batch_size is None:
            raise RuntimeError(
                "Reservoir state is not initialised. "
                "Pass batch_size= to initialise it before randomising."
            )
        self.reset_state(batch_size=batch_size)
        # ``reset_state(batch_size=...)`` always assigns a real tensor here
        # (batch_size is non-None past the guard above); narrows for mypy.
        assert self.state is not None
        # ``reset_state`` honours device/dtype of existing state or cell;
        # apply overrides if explicitly requested.
        if device is not None or dtype is not None:
            self.state = self.state.to(device=device, dtype=dtype)
    # State is non-None on both branches: either just initialised above or
    # already non-None at method entry.
    assert self.state is not None
    self.state = torch.randn_like(self.state)

ESNLayer

ESNLayer(reservoir_size: int, feedback_size: int, input_size: int | None = None, spectral_radius: float | None = None, bias: bool = True, bias_scaling: float = 1.0, activation: str = 'tanh', leak_rate: float | Tensor | Sequence[float] = 1.0, noise: float = 0.0, trainable: bool = False, feedback_initializer: InitializerSpec = None, input_initializer: InitializerSpec = None, topology: TopologySpec = None, seed: SeedLike = None)

Bases: BaseReservoirLayer

Stateful RNN reservoir layer for Echo State Networks.

Public-facing class. Internally creates an :class:ESNCell that owns all parameters and performs the single-step update; the sequence loop and state management live in the parent :class:BaseReservoirLayer.

PARAMETER	DESCRIPTION
reservoir_size	Number of reservoir units (hidden state dimension). TYPE: int
feedback_size	Dimension of feedback signal. Required for all reservoirs. TYPE: int
input_size	Dimension of driving inputs. If None, no driving input is expected. TYPE: int DEFAULT: None
spectral_radius	Target spectral radius for recurrent weights. Controls the "memory" and stability of the reservoir. If None, no spectral radius scaling is applied. TYPE: float DEFAULT: None
bias	Whether to include a bias term. TYPE: bool DEFAULT: True
bias_scaling	Scale of the random bias: entries are drawn from uniform(-bias_scaling, bias_scaling). Set to 0.0 for a zero bias (pre-0.5 behaviour). Ignored when bias=False. TYPE: float DEFAULT: 1.0
activation	Activation function for reservoir dynamics. TYPE: (tanh, relu, identity, sigmoid) DEFAULT: 'tanh'
leak_rate	Leaky integration rate in (0, 1]. A scalar 1.0 means no leaking (standard RNN update); smaller values create slower dynamics. A 1-D tensor/sequence of length reservoir_size assigns a per-neuron leak (heterogeneous timescales in one reservoir, matching RC.jl's leak_coefficient vector); every entry must be finite and in (0, 1]. See :class:~resdag.layers.cells.esn_cell.ESNCell for the full contract (a vector leak is stored as a buffer that moves with .to() / .double() and round-trips through save/load). TYPE: float or 1-D tensor/sequence of shape ``(reservoir_size,)`` DEFAULT: 1.0
noise	Standard deviation of additive Gaussian state noise injected after the activation (the classical ESN readout-conditioning regularizer). Active only in training mode (like dropout) and a no-op under :meth:~torch.nn.Module.eval. The default 0.0 disables it, leaving outputs bit-identical to the noiseless layer. Reproducible under the layer's seed. Must be non-negative. TYPE: float DEFAULT: 0.0
trainable	If True, reservoir weights are trainable via backpropagation. Standard ESNs use frozen (non-trainable) weights. TYPE: bool DEFAULT: False
feedback_initializer	Initializer for the feedback weight matrix. Accepts a registry name, (name, params), any matrix-building callable fn(rows, cols, **kw) -> matrix (or in-place fn(tensor), e.g. torch.nn.init.xavier_uniform_), (callable, params), or a configured initializer object. TYPE: str, callable, tuple, or InputFeedbackInitializer DEFAULT: None
input_initializer	Initializer for the input weight matrix. Same formats as feedback_initializer. Only used when input_size is provided. TYPE: str, callable, tuple, or InputFeedbackInitializer DEFAULT: None
topology	Structure of the recurrent weight matrix. Accepts a registry name (graph or matrix topology), (name, params), any matrix-building callable fn(n, **kw) -> matrix | nx.Graph (or in-place fn(tensor)), (callable, params), or a configured topology object. TYPE: str, callable, tuple, or TopologyInitializer DEFAULT: None
seed	Reproducibility seed that deterministically fixes the entire reservoir — topology (recurrent) matrix, feedback weights, input weights, and bias — so two layers built with the same seed are identical down to the last entry. Accepts a plain int or a :class:torch.Generator (handy for threading a per-trial generator from an HPO loop). This covers both explicit initializers/topologies and the default uniform(-1, 1) draws used when none is given. With seed set, a string-form topology='erdos_renyi' is reproducible without the ('erdos_renyi', {'seed': ...}) tuple form; an explicit seed in a tuple/object spec always wins. When seed=None (the default), the reservoir is still reproducible under torch.manual_seed because every generator (NumPy and torch) is derived from torch's global RNG. TYPE: int or Generator DEFAULT: None

ATTRIBUTE	DESCRIPTION
state	Current reservoir state of shape (batch, reservoir_size). None if not yet initialized. TYPE: Tensor or None

Examples:

Basic feedback-only reservoir:

>>> reservoir = ESNLayer(reservoir_size=500, feedback_size=10)
>>> feedback = torch.randn(4, 50, 10)  # (batch, time, features)
>>> output = reservoir(feedback)
>>> print(output.shape)
torch.Size([4, 50, 500])

Reservoir with driving input:

>>> reservoir = ESNLayer(
...     reservoir_size=500,
...     feedback_size=10,
...     input_size=5,
...     spectral_radius=0.95
... )
>>> feedback = torch.randn(4, 50, 10)
>>> driving = torch.randn(4, 50, 5)
>>> output = reservoir(feedback, driving)

Using graph topology by name:

>>> reservoir = ESNLayer(
...     reservoir_size=500,
...     feedback_size=10,
...     topology="erdos_renyi",
...     spectral_radius=0.9
... )

Using topology with custom parameters (tuple format):

>>> reservoir = ESNLayer(
...     reservoir_size=500,
...     feedback_size=10,
...     topology=("watts_strogatz", {"k": 6, "p": 0.3}),
...     feedback_initializer=("pseudo_diagonal", {"input_scaling": 0.5}),
...     spectral_radius=0.95
... )

Fully reproducible reservoir via the seed argument — recurrent matrix, feedback/input weights, and bias all match:

>>> a = ESNLayer(50, feedback_size=3, topology="erdos_renyi", seed=42)
>>> b = ESNLayer(50, feedback_size=3, topology="erdos_renyi", seed=42)
>>> torch.equal(a.weight_hh, b.weight_hh)
True
>>> torch.equal(a.weight_feedback, b.weight_feedback)
True
>>> torch.equal(a.bias_h, b.bias_h)
True

A :class:torch.Generator works too (e.g. an HPO per-trial generator):

>>> g1 = torch.Generator().manual_seed(7)
>>> g2 = torch.Generator().manual_seed(7)
>>> a = ESNLayer(50, feedback_size=3, seed=g1)
>>> b = ESNLayer(50, feedback_size=3, seed=g2)
>>> torch.equal(a.weight_feedback, b.weight_feedback)
True

The reservoir is also reproducible under torch.manual_seed without an explicit seed:

>>> torch.manual_seed(0)
>>> a = ESNLayer(50, feedback_size=3, topology="erdos_renyi")
>>> torch.manual_seed(0)
>>> b = ESNLayer(50, feedback_size=3, topology="erdos_renyi")
>>> torch.equal(a.weight_hh, b.weight_hh)
True

Stateful processing across batches:

>>> out1 = reservoir(data1)  # State initialized
>>> out2 = reservoir(data2)  # State carries over
>>> reservoir.reset_state()  # Manual reset
>>> out3 = reservoir(data3)  # Fresh state

See Also

resdag.layers.cells.esn_cell.ESNCell : Underlying single-step cell. resdag.layers.reservoirs.base_reservoir.BaseReservoirLayer : Base providing state management. resdag.init.topology.TopologyInitializer : Base class for topology init. resdag.init.input_feedback.InputFeedbackInitializer : Input init base. resdag.core.ESNModel : Model composition using reservoir layers.

Source code in src/resdag/layers/reservoirs/esn.py

def __init__(
    self,
    reservoir_size: int,
    feedback_size: int,
    input_size: int | None = None,
    spectral_radius: float | None = None,
    bias: bool = True,
    bias_scaling: float = 1.0,
    activation: str = "tanh",
    leak_rate: float | Tensor | Sequence[float] = 1.0,
    noise: float = 0.0,
    trainable: bool = False,
    feedback_initializer: InitializerSpec = None,
    input_initializer: InitializerSpec = None,
    topology: TopologySpec = None,
    seed: SeedLike = None,
) -> None:
    cell = ESNCell(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        input_size=input_size,
        spectral_radius=spectral_radius,
        bias=bias,
        bias_scaling=bias_scaling,
        activation=activation,
        leak_rate=leak_rate,
        noise=noise,
        trainable=trainable,
        feedback_initializer=feedback_initializer,
        input_initializer=input_initializer,
        topology=topology,
        seed=seed,
    )
    super().__init__(cell)

NGReservoir

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

Bases: NGReservoirLayer

Deprecated alias for :class:NGReservoirLayer.

Retained for backward compatibility. Instantiating this class emits a :class:DeprecationWarning and constructs an :class:NGReservoirLayer; the two are otherwise identical (NGReservoir is a subclass, so isinstance(NGReservoir(...), NGReservoirLayer) holds).

.. deprecated:: Use :class:NGReservoirLayer instead. NGReservoir will be removed in a future release.

See Also

NGReservoirLayer : The canonical class this aliases.

Source code in src/resdag/layers/reservoirs/ngrc.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    warnings.warn(
        "NGReservoir is deprecated and will be removed in a future release; "
        "use NGReservoirLayer instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    super().__init__(*args, **kwargs)

readouts

Readout Layers

This module provides readout layer implementations for ESN models.

CLASS	DESCRIPTION
ReadoutLayer	Base per-timestep linear layer with fitting interface.
CGReadoutLayer	Readout with Conjugate Gradient ridge regression solver.
RidgeReadoutLayer	Direct ridge readout via Cholesky (solver='cholesky') or LU (solver='solve') factorisation of the normal equations.
CholeskyReadoutLayer	Single-shot Cholesky ridge readout; the streaming-path direct solver.
SVDReadoutLayer	Readout solved via SVD with Tikhonov filter factors; robust to rank-deficient Gram matrices and alpha=0.
PinvReadoutLayer	Least-squares readout via torch.linalg.lstsq / pinv with an rcond cutoff.
IncrementalRidgeReadoutLayer	Streaming ridge readout: partial_fit accumulates sufficient statistics chunk-by-chunk, finalize solves once. For the DataLoader / long-sequence path.
RLSReadout	True-online Recursive Least Squares readout with a forgetting factor: partial_fit refines the weights one sample at a time from a running inverse-correlation matrix. Tracks non-stationary targets (lambda < 1) and converges to ordinary least squares (lambda = 1).
FISTAReadoutLayer	Sparse / robust readout: L1 (lasso) or elastic-net regression solved with FISTA (accelerated proximal gradient), batched over all outputs.
The old names ``IncrementalRidgeReadout`` and ``FISTAReadout`` are retained as
deprecated aliases of ``IncrementalRidgeReadoutLayer`` / ``FISTAReadoutLayer``.

Examples:

>>> from resdag.layers.readouts import CGReadoutLayer
>>> readout = CGReadoutLayer(
...     in_features=200,
...     out_features=3,
...     alpha=1e-6,
...     name="output",
... )
>>> # Fit using ESNTrainer or directly
>>> readout.fit(states, targets)
>>> output = readout(states)

See Also

resdag.training.ESNTrainer : Trainer that uses these readouts. resdag.layers.ESNLayer : ESN layer for generating states.

ReadoutLayer

ReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False)

Bases: Linear

Per-timestep linear layer with custom fitting for ESN training.

This layer extends :class:torch.nn.Linear with:

• Per-timestep application to sequence tensors (B, T, F)
• Named identification for multi-readout architectures
• Custom fit() interface for classical ESN training

The layer applies the same linear transformation independently to each timestep in a sequence:

.. code-block:: text

Input:  (B, T, F_in)  -> Reshape to (B*T, F_in)
Apply:  linear(x) = x @ W.T + b
Output: (B*T, F_out) -> Reshape to (B, T, F_out)

This matches classical ESN semantics where readouts are fitted across the entire sequence at once using ridge regression.

PARAMETER	DESCRIPTION
in_features	Size of input features. TYPE: int
out_features	Size of output features. TYPE: int
bias	Whether to include a bias term. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. If False, weights are frozen (standard ESN behavior). TYPE: bool DEFAULT: False

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
name	Name of this readout layer. TYPE: str or None
is_fitted	True if fit() has been called successfully. TYPE: bool

Examples:

Basic usage:

>>> readout = ReadoutLayer(in_features=100, out_features=10)
>>> x = torch.randn(2, 20, 100)  # (batch, seq_len, features)
>>> y = readout(x)
>>> print(y.shape)
torch.Size([2, 20, 10])

Named readout for multi-output architectures:

>>> readout1 = ReadoutLayer(100, 10, name="position")
>>> readout2 = ReadoutLayer(100, 3, name="velocity")

See Also

CGReadoutLayer : Readout with Conjugate Gradient solver. resdag.training.ESNTrainer : Trainer for fitting readouts.

Source code in src/resdag/layers/readouts/base.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
) -> None:
    super().__init__(in_features, out_features, bias)

    self._name = name
    # ``is_fitted`` is stored as a registered (persistent) buffer so it
    # survives ``state_dict`` / ``load_state_dict`` round-trips. A plain
    # Python attribute would silently reset to ``False`` on reload, leaving
    # a model with fitted weights but ``is_fitted == False``.
    self.register_buffer("_is_fitted", torch.zeros((), dtype=torch.bool))
    # Plain Python mirror of the ``_is_fitted`` buffer. Reading the buffer
    # via ``.item()`` is a device->host sync on CUDA — cheap once, but
    # ruinous inside a per-step autoregressive forecast loop (and it also
    # graph-breaks ``torch.compile``). The mirror lets hot paths (e.g.
    # ``IncrementalRidgeReadoutLayer.forward``) check fittedness with a plain
    # attribute read and no sync. It is kept coherent with the buffer via
    # :meth:`_set_fitted` and re-derived from the buffer after every
    # ``load_state_dict`` (see :meth:`_load_from_state_dict`).
    self._fitted_flag: bool = False
    # Register a post-load hook so subclasses/instances loaded via
    # ``load_state_dict`` re-sync the Python mirror from the restored buffer.
    self.register_load_state_dict_post_hook(_sync_fitted_flag_post_hook)

    # ``trainable`` is a live property whose setter flips ``requires_grad_``
    # on every parameter (``weight``, ``bias``).  ``nn.Linear.__init__`` has
    # already created those parameters, so a single assignment both records
    # the flag and applies the freeze/unfreeze — no separate
    # ``_freeze_weights`` step is needed.
    self.trainable = trainable

trainable property writable

trainable: bool

bool : Whether the readout weights participate in backpropagation.

A live switch: assigning to it flips requires_grad_ on every readout parameter (weight and, when present, bias). A post-construction readout.trainable = True therefore genuinely unfreezes the readout for a gradient-based rollout rather than silently shadowing a construction-time-only flag. The getter reflects the current state.

name property

name: str | None

str or None : Name of this readout layer.

is_fitted property

is_fitted: bool

bool : True if fit() has been called successfully.

Backed by a registered buffer (so the flag is preserved across state_dict / load_state_dict round-trips) but read from a plain Python mirror kept coherent with it. Reading the mirror avoids the buffer .item() device->host sync — negligible once, but ruinous inside a per-step autoregressive forecast loop and a torch.compile graph-break.

forward

forward(input: Tensor) -> Tensor

Apply linear transformation to input.

Handles both 2D (batch, features) and 3D (batch, seq_len, features) inputs. For 3D inputs, applies the linear transformation independently to each timestep.

PARAMETER	DESCRIPTION
input	Input tensor of shape (B, F) or (B, T, F). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Output tensor of shape (B, F_out) or (B, T, F_out).

RAISES	DESCRIPTION
ValueError	If input has neither 2 nor 3 dimensions.

Examples:

>>> readout = ReadoutLayer(100, 10)
>>> x_2d = torch.randn(4, 100)
>>> y_2d = readout(x_2d)  # (4, 10)
>>> x_3d = torch.randn(4, 50, 100)
>>> y_3d = readout(x_3d)  # (4, 50, 10)

Source code in src/resdag/layers/readouts/base.py

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """
    Apply linear transformation to input.

    Handles both 2D ``(batch, features)`` and 3D ``(batch, seq_len, features)``
    inputs. For 3D inputs, applies the linear transformation independently
    to each timestep.

    Parameters
    ----------
    input : torch.Tensor
        Input tensor of shape ``(B, F)`` or ``(B, T, F)``.

    Returns
    -------
    torch.Tensor
        Output tensor of shape ``(B, F_out)`` or ``(B, T, F_out)``.

    Raises
    ------
    ValueError
        If input has neither 2 nor 3 dimensions.

    Examples
    --------
    >>> readout = ReadoutLayer(100, 10)
    >>> x_2d = torch.randn(4, 100)
    >>> y_2d = readout(x_2d)  # (4, 10)
    >>> x_3d = torch.randn(4, 50, 100)
    >>> y_3d = readout(x_3d)  # (4, 50, 10)
    """
    if input.dim() == 2:
        return super().forward(input)

    elif input.dim() == 3:
        batch_size, seq_len, features = input.shape
        input_reshaped = input.reshape(batch_size * seq_len, features)
        output_reshaped = super().forward(input_reshaped)
        output = output_reshaped.reshape(batch_size, seq_len, self.out_features)
        return output

    else:
        raise ValueError(
            f"ReadoutLayer expects 2D (B, F) or 3D (B, T, F) input, "
            f"got {input.dim()}D tensor with shape {input.shape}"
        )

fit

fit(states: Tensor, targets: Tensor) -> None

Fit readout weights to the given (states, targets) pair.

This base implementation handles the bookkeeping that every readout needs: shape normalisation from (B, T, F) → (B*T, F) plus sample-count and out-features validation. The actual algebraic solve is delegated to :meth:_fit_impl, which subclasses override.

PARAMETER	DESCRIPTION
states	Input states of shape (B, T, F_in) or (N, F_in). TYPE: Tensor
targets	Target outputs of shape (B, T, F_out) or (N, F_out). TYPE: Tensor

RAISES	DESCRIPTION
NotImplementedError	The base ReadoutLayer does not implement an algebraic solver; use a subclass such as :class:CGReadoutLayer.
ValueError	If states and targets disagree on the sample dimension after flattening, if the state's feature dimension does not match self.in_features, if the target's feature dimension does not match self.out_features, or if states or targets contain non-finite (NaN or Inf) values.

Notes

After fit() returns successfully, :attr:is_fitted is True.

See Also

CGReadoutLayer : Concrete implementation using Conjugate Gradient.

Source code in src/resdag/layers/readouts/base.py

@torch.no_grad()
def fit(
    self,
    states: torch.Tensor,
    targets: torch.Tensor,
) -> None:
    """
    Fit readout weights to the given (states, targets) pair.

    This base implementation handles the bookkeeping that every readout
    needs: shape normalisation from ``(B, T, F)`` → ``(B*T, F)`` plus
    sample-count and out-features validation.  The actual algebraic
    solve is delegated to :meth:`_fit_impl`, which subclasses override.

    Parameters
    ----------
    states : torch.Tensor
        Input states of shape ``(B, T, F_in)`` or ``(N, F_in)``.
    targets : torch.Tensor
        Target outputs of shape ``(B, T, F_out)`` or ``(N, F_out)``.

    Raises
    ------
    NotImplementedError
        The base ``ReadoutLayer`` does not implement an algebraic solver;
        use a subclass such as :class:`CGReadoutLayer`.
    ValueError
        If ``states`` and ``targets`` disagree on the sample dimension
        after flattening, if the state's feature dimension does not match
        ``self.in_features``, if the target's feature dimension does not
        match ``self.out_features``, or if ``states`` or ``targets``
        contain non-finite (NaN or Inf) values.

    Notes
    -----
    After ``fit()`` returns successfully, :attr:`is_fitted` is ``True``.

    See Also
    --------
    CGReadoutLayer : Concrete implementation using Conjugate Gradient.
    """
    if states.dim() == 3:
        b, t, f = states.shape
        states = states.reshape(b * t, f)
    if targets.dim() == 3:
        b, t, f = targets.shape
        targets = targets.reshape(b * t, f)

    readout_id = f"'{self._name}'" if self._name is not None else type(self).__name__
    if states.shape[0] != targets.shape[0]:
        raise ValueError(
            f"{type(self).__name__}.fit({readout_id}): sample count mismatch. "
            f"States have {states.shape[0]} samples after flattening, "
            f"targets have {targets.shape[0]}."
        )
    if states.shape[1] != self.in_features:
        raise ValueError(
            f"{type(self).__name__}.fit({readout_id}): state feature dimension "
            f"({states.shape[1]}) does not match readout in_features "
            f"({self.in_features})."
        )
    if targets.shape[1] != self.out_features:
        raise ValueError(
            f"{type(self).__name__}.fit({readout_id}): target feature dimension "
            f"({targets.shape[1]}) does not match readout out_features "
            f"({self.out_features})."
        )

    # Reject non-finite inputs *before* the algebraic solve. A diverged
    # reservoir (or a corrupt target) feeds NaN/Inf into the Gram-matrix
    # matmuls, which would otherwise propagate straight into the fitted
    # weights and leave the layer ``is_fitted`` with NaN parameters and no
    # error. Failing loudly here is far easier to diagnose.
    if not torch.isfinite(states).all():
        raise ValueError(
            f"{type(self).__name__}.fit({readout_id}): states contain "
            f"non-finite values (NaN or Inf). The readout solve requires "
            f"finite inputs — clean or impute the reservoir states (a "
            f"diverged reservoir is the usual cause) before fitting."
        )
    if not torch.isfinite(targets).all():
        raise ValueError(
            f"{type(self).__name__}.fit({readout_id}): targets contain "
            f"non-finite values (NaN or Inf). The readout solve requires "
            f"finite inputs — clean or impute the targets before fitting."
        )

    coefs, intercept = self._fit_impl(states, targets)

    self.weight.copy_(coefs.T.to(self.weight.dtype))
    if self.bias is not None and intercept is not None:
        self.bias.copy_(intercept.to(self.bias.dtype))

    self._set_fitted(True)

RidgeReadoutLayer

RidgeReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False, alpha: float = 1e-06, solver: str = 'cholesky', use_float64: bool = True, gram_dtype: dtype | None = None)

Bases: ReadoutLayer

Readout layer with a direct ridge-regression solver.

For the readout widths typical of reservoir computing (in_features = 100–2000) the (F, F) Gram matrix factorises exactly and almost instantly. A direct Cholesky solve is both faster and exact (up to floating point), whereas :class:CGReadoutLayer only approximates the solution and can under-converge on the ill-conditioned concatenated-input readouts produced by ott_esn / power_augmented.

The solver minimises the same ridge objective as :class:CGReadoutLayer,

.. math::

\lVert X W - Y \rVert_2^2 + \alpha \lVert W \rVert_2^2 ,

by forming the regularised normal equations (XᵀX + αI) W = Xᵀy and solving them with one of:

• solver='cholesky' (default) — torch.linalg.cholesky followed by torch.linalg.cholesky_solve. Exploits the symmetric positive definiteness of XᵀX + αI (guaranteed for alpha > 0) and is the fastest, most numerically stable option for well-conditioned fits.
• solver='solve' — torch.linalg.solve (LU). Slightly slower but does not require positive definiteness, so it tolerates alpha == 0 on a full-rank Gram.

Solver-selection guide

• Default to RidgeReadoutLayer(solver='cholesky') for any well-conditioned readout with alpha > 0 and F <= 2000. It is the fastest exact option and matches :class:CGReadoutLayer to < 1e-8.
• Use solver='solve' when you need alpha == 0 (pure least squares) and the Gram is full rank but you don't want the SVD machinery.
• If the Gram is rank deficient (high-degree NG-RC feature maps, more features than samples), a Cholesky/LU solve will fail or be unstable — reach for :class:SVDReadoutLayer (filter factors) or :class:PinvReadoutLayer (lstsq / pinv) instead.

PARAMETER	DESCRIPTION
in_features	Size of input features (reservoir state dimension). TYPE: int
out_features	Size of output features (prediction dimension). TYPE: int
bias	Whether to include a bias term. When True the data is centered and an unpenalised intercept is recovered; when False the raw normal equations are solved with no intercept. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and as the key into targets by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. TYPE: bool DEFAULT: False
alpha	L2 regularization strength. Must be non-negative. alpha == 0 is only safe with solver='solve' on a full-rank Gram. TYPE: float DEFAULT: 1e-6
solver	Direct solver to use for the normal equations. TYPE: (cholesky, solve) DEFAULT: 'cholesky'
use_float64	If True the small (F, F) factorisation runs in float64 and the result is cast back to the parameter dtype. TYPE: bool DEFAULT: True
gram_dtype	Dtype for forming the Gram matrix and right-hand side (the heavy (N, F) matmuls). Default None is automatic: float64 on CPU (cheap there), the input dtype on CUDA (float64 matmuls are crippled on consumer GPUs). Pass torch.float64 to force full-precision Gram formation on any device. TYPE: dtype DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
alpha	L2 regularization strength. TYPE: float
solver	The selected direct solver. TYPE: str

Examples:

>>> readout = RidgeReadoutLayer(in_features=200, out_features=3, alpha=1e-6)
>>> states = torch.randn(8, 50, 200)  # (batch, time, features)
>>> targets = torch.randn(8, 50, 3)
>>> readout.fit(states, targets)
>>> output = readout(states)
>>> output.shape
torch.Size([8, 50, 3])

See Also

CGReadoutLayer : Iterative Conjugate Gradient ridge readout. SVDReadoutLayer : SVD filter-factor readout for rank-deficient Gram. PinvReadoutLayer : lstsq / pseudo-inverse readout. resdag.training.ESNTrainer : Trainer that uses this for fitting.

Source code in src/resdag/layers/readouts/ridge_readout.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
    alpha: float = 1e-6,
    solver: str = "cholesky",
    use_float64: bool = True,
    gram_dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias, name, trainable)
    if alpha < 0:
        raise ValueError(f"alpha must be non-negative, got {alpha}.")
    if solver not in self._VALID_SOLVERS:
        raise ValueError(f"solver must be one of {self._VALID_SOLVERS}, got {solver!r}.")
    self.alpha = alpha
    self.solver = solver
    self.use_float64 = use_float64
    self.gram_dtype = gram_dtype

CholeskyReadoutLayer

CholeskyReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False, alpha: float = 1e-06, use_float64: bool = True, gram_dtype: dtype | None = None)

Bases: ReadoutLayer

Readout layer that solves ridge regression via a Cholesky factorisation.

The regularised normal equations (XᵀX + αI) W = Xᵀy are formed (with analytic centering when a bias is fitted) and solved with :func:torch.linalg.cholesky followed by :func:torch.cholesky_solve. The (F, F) system is symmetric positive definite for alpha > 0, so the Cholesky factorisation is the fastest, most numerically stable exact solve for the readout widths typical of reservoir computing (in_features = 100–2000).

Unlike :class:CGReadoutLayer, which only approximates the solution iteratively, this layer returns the exact ridge solution (up to floating point) in a single shot. It is the single-batch counterpart of the streaming :class:IncrementalRidgeReadoutLayer: both ultimately Cholesky-solve the same regularised Gram system, so a full-batch CholeskyReadoutLayer fit and an accumulated IncrementalRidgeReadoutLayer fit over the same data agree to within floating-point tolerance.

The solver minimises the same ridge objective as :class:CGReadoutLayer,

.. math::

\lVert X W - Y \rVert_2^2 + \alpha \lVert W \rVert_2^2 .

PARAMETER	DESCRIPTION
in_features	Size of input features (reservoir state dimension). TYPE: int
out_features	Size of output features (prediction dimension). TYPE: int
bias	Whether to include a bias term. When True the data is centered and an unpenalised intercept is recovered; when False the raw normal equations are solved with no intercept. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and as the key into targets by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. TYPE: bool DEFAULT: False
alpha	L2 regularization strength. Must be non-negative. alpha > 0 guarantees the Gram is positive definite; with alpha == 0 the factorisation succeeds only on a full-rank Gram (reach for :class:SVDReadoutLayer if the Gram may be rank deficient). TYPE: float DEFAULT: 1e-6
use_float64	If True the small (F, F) factorisation runs in float64 and the result is cast back to the parameter dtype. TYPE: bool DEFAULT: True
gram_dtype	Dtype for forming the Gram matrix and right-hand side (the heavy (N, F) matmuls). Default None is automatic: float64 on CPU (cheap there), the input dtype on CUDA (float64 matmuls are crippled on consumer GPUs). Pass torch.float64 to force full-precision Gram formation on any device. TYPE: dtype DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
alpha	L2 regularization strength. TYPE: float

Examples:

>>> readout = CholeskyReadoutLayer(in_features=200, out_features=3, alpha=1e-6)
>>> states = torch.randn(8, 50, 200)  # (batch, time, features)
>>> targets = torch.randn(8, 50, 3)
>>> readout.fit(states, targets)
>>> output = readout(states)
>>> output.shape
torch.Size([8, 50, 3])

See Also

CGReadoutLayer : Iterative Conjugate Gradient ridge readout. RidgeReadoutLayer : Direct ridge readout with a solver switch. IncrementalRidgeReadoutLayer : Streaming partial_fit / finalize counterpart. resdag.training.ESNTrainer : Trainer that uses this for fitting.

Source code in src/resdag/layers/readouts/cholesky_readout.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
    alpha: float = 1e-6,
    use_float64: bool = True,
    gram_dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias, name, trainable)
    if alpha < 0:
        raise ValueError(f"alpha must be non-negative, got {alpha}.")
    self.alpha = alpha
    self.use_float64 = use_float64
    self.gram_dtype = gram_dtype

CGReadoutLayer

CGReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False, max_iter: int = 100, tol: float = 1e-05, alpha: float = 1e-06, use_float64: bool = True, gram_dtype: dtype | None = None)

Bases: ReadoutLayer

Readout layer with Conjugate Gradient ridge regression solver.

This layer extends :class:ReadoutLayer with an efficient Conjugate Gradient (CG) solver for fitting weights via ridge regression. The CG solver is:

• Memory efficient (doesn't form full normal equations matrix)
• GPU accelerated
• Numerically stable (uses float64 internally)
• Supports batched time-series data

The solver finds weights W that minimize:

.. math::

||XW - Y||^2 + \alpha ||W||^2

where :math:\alpha is the regularization strength.

PARAMETER	DESCRIPTION
in_features	Size of input features (reservoir state dimension). TYPE: int
out_features	Size of output features (prediction dimension). TYPE: int
bias	Whether to include a bias term. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. TYPE: bool DEFAULT: False
alpha	L2 regularization strength. Must be non-negative. Larger values provide more regularization (smoother outputs, less overfitting). TYPE: float DEFAULT: 1e-6
max_iter	Maximum number of CG iterations. Must be a positive integer. TYPE: int DEFAULT: 100
tol	Convergence tolerance for the CG solver. Must be positive. Each output column is considered converged when its residual-norm-squared falls below tol**2. TYPE: float DEFAULT: 1e-5
use_float64	If True (default), the CG iterations on the small (in_features, in_features) system run in float64 and the result is cast back to the layer's parameter dtype. Cheap on every device — the expensive (N, F) Gram-formation matmuls stay in the input dtype regardless (see gram_dtype). TYPE: bool DEFAULT: True
gram_dtype	Dtype for forming the Gram matrix and right-hand side (the heavy (N, F) matmuls). Default None is automatic: float64 on CPU (cheap there), the input dtype on CUDA — float64 matmuls run at 1/32–1/64 speed on consumer GPUs and are the classic reason ESN training measures slower on GPU than CPU. Pass torch.float64 to force full-precision Gram formation on any device (needed only for badly scaled states, e.g. unnormalized inputs concatenated into the readout; prefer normalizing the data). TYPE: dtype DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
alpha	L2 regularization strength. TYPE: float
max_iter	Maximum CG iterations. TYPE: int
tol	Convergence tolerance. TYPE: float

Examples:

Basic usage:

>>> readout = CGReadoutLayer(in_features=100, out_features=10, alpha=1e-6)
>>> states = torch.randn(32, 50, 100)  # (batch, time, features)
>>> targets = torch.randn(32, 50, 10)
>>> readout.fit(states, targets)
>>> output = readout(states)
>>> print(output.shape)
torch.Size([32, 50, 10])

With custom regularization:

>>> readout = CGReadoutLayer(100, 10, alpha=1e-4)  # Stronger regularization
>>> readout.fit(states, targets)

See Also

ReadoutLayer : Base readout layer class. resdag.training.ESNTrainer : Trainer that uses this for fitting.

Source code in src/resdag/layers/readouts/cg_readout.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
    max_iter: int = 100,
    tol: float = 1e-5,
    alpha: float = 1e-6,
    use_float64: bool = True,
    gram_dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias, name, trainable)
    if alpha < 0:
        raise ValueError(f"alpha must be non-negative, got {alpha}.")
    if tol <= 0:
        raise ValueError(f"tol must be positive, got {tol}.")
    if max_iter < 1:
        raise ValueError(f"max_iter must be a positive integer, got {max_iter}.")
    self.max_iter = max_iter
    self.tol = tol
    self.alpha = alpha
    self.use_float64 = use_float64
    self.gram_dtype = gram_dtype

SVDReadoutLayer

SVDReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False, alpha: float = 0.0, rcond: float = 1e-15, use_float64: bool = True, gram_dtype: dtype | None = None)

Bases: ReadoutLayer

Readout layer solved via SVD with Tikhonov filter factors.

The (optionally centered) design matrix X is decomposed as X = U diag(s) Vᵀ and the ridge solution is reconstructed with filter factors

.. math::

W = V \, \operatorname{diag}\!\left(\frac{s}{s^2 + \alpha}\right)
    U^\top Y .

For alpha > 0 this is exactly the ridge solution; for alpha == 0 it collapses to the minimum-norm least-squares solution, with singular values at or below an rcond cutoff dropped so a rank-deficient Gram (more features than independent samples, or the near-collinear columns produced by high-degree NG-RC feature maps) is handled gracefully instead of blowing up the way a Cholesky / normal-equations solve would.

Working from the SVD of X directly (rather than the Gram XᵀX) squares the effective condition number's square root away, so the small singular values that an iterative CG solver under-resolves are recovered faithfully — giving a strictly lower train residual on the ill-conditioned fits where CG stalls.

Solver-selection guide

• Reach for :class:SVDReadoutLayer when the Gram is rank deficient or severely ill-conditioned: alpha == 0 least squares, high-degree NG-RC feature maps, or wide concatenated-input readouts where CG under-converges. It is the most robust of the readouts.
• For a well-conditioned readout with alpha > 0 and F <= 2000, :class:RidgeReadoutLayer (solver='cholesky') is faster and equally accurate — prefer it there.
• :class:PinvReadoutLayer is the closely-related lstsq / pinv route; SVD here gives explicit control over the Tikhonov filter and rcond.

PARAMETER	DESCRIPTION
in_features	Size of input features (reservoir state dimension). TYPE: int
out_features	Size of output features (prediction dimension). TYPE: int
bias	Whether to include a bias term. When True the data is centered and an unpenalised intercept is recovered; when False the raw design matrix is decomposed with no intercept. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and as the key into targets by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. TYPE: bool DEFAULT: False
alpha	Tikhonov (L2) regularization strength. Must be non-negative. 0 is fully supported and yields the minimum-norm least-squares solution. TYPE: float DEFAULT: 0.0
rcond	Relative cutoff for small singular values. Singular values at or below rcond * s_max are treated as zero (their filter factor is set to zero), which is what makes rank-deficient problems well posed. Only active for alpha == 0; for alpha > 0 the Tikhonov filter already damps small singular values. TYPE: float DEFAULT: 1e-15
use_float64	If True the SVD runs in float64 and the result is cast back to the parameter dtype. TYPE: bool DEFAULT: True
gram_dtype	Dtype for forming / casting the design matrix. Default None is automatic: float64 on CPU, the input dtype on CUDA. Pass torch.float64 to force full precision on any device. TYPE: dtype DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
alpha	Tikhonov regularization strength. TYPE: float
rcond	Relative singular-value cutoff. TYPE: float

Examples:

>>> readout = SVDReadoutLayer(in_features=200, out_features=3, alpha=0.0)
>>> states = torch.randn(8, 50, 200)  # (batch, time, features)
>>> targets = torch.randn(8, 50, 3)
>>> readout.fit(states, targets)
>>> output = readout(states)
>>> output.shape
torch.Size([8, 50, 3])

See Also

RidgeReadoutLayer : Cholesky / solve direct readout (faster, full rank). PinvReadoutLayer : lstsq / pseudo-inverse readout. CGReadoutLayer : Iterative Conjugate Gradient ridge readout. resdag.training.ESNTrainer : Trainer that uses this for fitting.

Source code in src/resdag/layers/readouts/svd_readout.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
    alpha: float = 0.0,
    rcond: float = 1e-15,
    use_float64: bool = True,
    gram_dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias, name, trainable)
    if alpha < 0:
        raise ValueError(f"alpha must be non-negative, got {alpha}.")
    if rcond < 0:
        raise ValueError(f"rcond must be non-negative, got {rcond}.")
    self.alpha = alpha
    self.rcond = rcond
    self.use_float64 = use_float64
    self.gram_dtype = gram_dtype

PinvReadoutLayer

PinvReadoutLayer(in_features: int, out_features: int, bias: bool = True, name: str | None = None, trainable: bool = False, rcond: float = 1e-15, solver: str = 'lstsq', use_float64: bool = True, gram_dtype: dtype | None = None)

Bases: ReadoutLayer

Readout layer solved via least squares / pseudo-inverse.

Solves the (optionally centered) least-squares problem min_W ||X W - Y|| directly, without forming the Gram matrix, using one of:

• solver='lstsq' (default) — torch.linalg.lstsq with the gelsd driver, a divide-and-conquer SVD-based least-squares solver that returns the minimum-norm solution for rank-deficient X and honours the rcond cutoff.
• solver='pinv' — form the Moore–Penrose pseudo-inverse torch.linalg.pinv(X, rcond=...) and apply it to Y. Equivalent in result but materialises the (F, N) pseudo-inverse explicitly.

This is a pure least-squares readout: there is no Tikhonov penalty (alpha). Regularisation comes only from the rcond truncation of small singular values. If you need an explicit ridge penalty use :class:RidgeReadoutLayer or :class:SVDReadoutLayer (which exposes both alpha and rcond).

Solver-selection guide

• Reach for :class:PinvReadoutLayer for unregularised least squares on a possibly rank-deficient design matrix when you want the minimum-norm solution and don't need a Tikhonov alpha.
• solver='lstsq' (default) is the right choice almost always — it never materialises the pseudo-inverse. Use solver='pinv' only if you specifically want the pseudo-inverse matrix itself.
• For an explicit ridge penalty, prefer :class:SVDReadoutLayer (filter factors, alpha + rcond) or :class:RidgeReadoutLayer (Cholesky, fast, full rank).

PARAMETER	DESCRIPTION
in_features	Size of input features (reservoir state dimension). TYPE: int
out_features	Size of output features (prediction dimension). TYPE: int
bias	Whether to include a bias term. When True the data is centered and an unpenalised intercept is recovered; when False the raw design matrix is used with no intercept. TYPE: bool DEFAULT: True
name	Name for this readout layer. Used for identification in multi-readout architectures and as the key into targets by :class:ESNTrainer. TYPE: str DEFAULT: None
trainable	If True, weights are trainable via backpropagation. TYPE: bool DEFAULT: False
rcond	Relative cutoff for small singular values. Singular values at or below rcond * s_max are dropped, making rank-deficient problems well posed. TYPE: float DEFAULT: 1e-15
solver	Whether to solve with torch.linalg.lstsq or an explicit torch.linalg.pinv. TYPE: (lstsq, pinv) DEFAULT: 'lstsq'
use_float64	If True the least-squares solve runs in float64 and the result is cast back to the parameter dtype. TYPE: bool DEFAULT: True
gram_dtype	Dtype for forming / casting the design matrix. Default None is automatic: float64 on CPU, the input dtype on CUDA. Pass torch.float64 to force full precision on any device. TYPE: dtype DEFAULT: None

ATTRIBUTE	DESCRIPTION
weight	Weight matrix of shape (out_features, in_features). TYPE: Parameter
bias	Bias vector of shape (out_features,), or None if bias=False. TYPE: Parameter or None
rcond	Relative singular-value cutoff. TYPE: float
solver	The selected least-squares solver. TYPE: str

Examples:

>>> readout = PinvReadoutLayer(in_features=200, out_features=3)
>>> states = torch.randn(8, 50, 200)  # (batch, time, features)
>>> targets = torch.randn(8, 50, 3)
>>> readout.fit(states, targets)
>>> output = readout(states)
>>> output.shape
torch.Size([8, 50, 3])

See Also

SVDReadoutLayer : SVD filter-factor readout with explicit alpha. RidgeReadoutLayer : Cholesky / solve direct ridge readout. CGReadoutLayer : Iterative Conjugate Gradient ridge readout. resdag.training.ESNTrainer : Trainer that uses this for fitting.

Source code in src/resdag/layers/readouts/pinv_readout.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    name: str | None = None,
    trainable: bool = False,
    rcond: float = 1e-15,
    solver: str = "lstsq",
    use_float64: bool = True,
    gram_dtype: torch.dtype | None = None,
) -> None:
    super().__init__(in_features, out_features, bias, name, trainable)
    if rcond < 0:
        raise ValueError(f"rcond must be non-negative, got {rcond}.")
    if solver not in self._VALID_SOLVERS:
        raise ValueError(f"solver must be one of {self._VALID_SOLVERS}, got {solver!r}.")
    self.rcond = rcond
    self.solver = solver
    self.use_float64 = use_float64
    self.gram_dtype = gram_dtype

IncrementalRidgeReadout

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

Bases: IncrementalRidgeReadoutLayer

Deprecated alias for :class:IncrementalRidgeReadoutLayer.

Retained for backward compatibility. Instantiating this class emits a :class:DeprecationWarning and constructs an :class:IncrementalRidgeReadoutLayer; the two are otherwise identical (IncrementalRidgeReadout is a subclass, so isinstance(IncrementalRidgeReadout(...), IncrementalRidgeReadoutLayer) holds).

.. deprecated:: Use :class:IncrementalRidgeReadoutLayer instead. IncrementalRidgeReadout will be removed in a future release.

See Also

IncrementalRidgeReadoutLayer : The canonical class this aliases.

Source code in src/resdag/layers/readouts/incremental_ridge.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    warnings.warn(
        "IncrementalRidgeReadout is deprecated and will be removed in a future "
        "release; use IncrementalRidgeReadoutLayer instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    super().__init__(*args, **kwargs)

transforms

Transform Layers

Specialized transformation layers for advanced ESN architectures, including state augmentation, feature manipulation, and regularization.

CLASS	DESCRIPTION
Concatenate	Concatenates multiple inputs along the feature dimension.
ExtendedSquare	Appends the elementwise square of every feature ([x, x**2]).
FeaturePartitioner	Partitions input features into overlapping groups.
NLAT2	T₂ nonlinear algebraic transform (x[i-1] * x[i-2] at even features).
NLAT3	T₃ nonlinear algebraic transform (x[i-1] * x[i+1] at interior even features).
Pad	Appends a constant column to the feature dimension ([x, padding]).
PartialSquare	Squares a leading fraction (or count) of features in place.
Power	Per-feature exponentiation by a fixed power.
SelectiveDropout	Per-feature dropout with selectivity control.
SelectiveExponentiation	Per-feature exponentiation transformation (squares even-indexed units; index=0, exponent=2 is the T₁/NLAT1 transform).
Standardize	Per-feature z-score standardization with buffer-stored mean/std, fit, and an exact inverse.

Examples:

>>> from resdag.layers.transforms import Concatenate, SelectiveExponentiation
>>> import torch
>>>
>>> concat = Concatenate()
>>> x1 = torch.randn(4, 100, 50)
>>> x2 = torch.randn(4, 100, 50)
>>> combined = concat(x1, x2)  # (4, 100, 100)

Concatenate

Bases: Module

Concatenate multiple tensors along the feature (last) dimension.

This layer takes any number of input tensors and concatenates them along the last dimension. Useful for combining outputs from parallel branches in multi-reservoir or ensemble architectures.

All input tensors must have the same shape except for the last dimension (features).

ATTRIBUTE	DESCRIPTION
None	This layer has no learnable parameters.

Examples:

Basic concatenation:

>>> import torch
>>> from resdag.layers.transforms import Concatenate
>>>
>>> concat = Concatenate()
>>> x1 = torch.randn(32, 50, 100)  # (batch, time, features1)
>>> x2 = torch.randn(32, 50, 200)  # (batch, time, features2)
>>> y = concat(x1, x2)
>>> print(y.shape)
torch.Size([32, 50, 300])

Multiple inputs:

>>> x3 = torch.randn(32, 50, 50)
>>> y = concat(x1, x2, x3)
>>> print(y.shape)
torch.Size([32, 50, 350])

See Also

FeaturePartitioner : Splits features into overlapping partitions.

forward

forward(*inputs: Tensor) -> Tensor

Concatenate input tensors along the last dimension.

PARAMETER	DESCRIPTION
*inputs	Variable number of input tensors to concatenate. All tensors must have the same shape except for the last dimension. TYPE: Tensor DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	Concatenated tensor with combined feature dimension.

RAISES	DESCRIPTION
RuntimeError	If tensors have incompatible shapes for concatenation.

Examples:

>>> concat = Concatenate()
>>> a = torch.randn(4, 10, 20)
>>> b = torch.randn(4, 10, 30)
>>> result = concat(a, b)
>>> print(result.shape)
torch.Size([4, 10, 50])

Source code in src/resdag/layers/transforms/concatenate.py

def forward(self, *inputs: torch.Tensor) -> torch.Tensor:
    """
    Concatenate input tensors along the last dimension.

    Parameters
    ----------
    *inputs : torch.Tensor
        Variable number of input tensors to concatenate. All tensors
        must have the same shape except for the last dimension.

    Returns
    -------
    torch.Tensor
        Concatenated tensor with combined feature dimension.

    Raises
    ------
    RuntimeError
        If tensors have incompatible shapes for concatenation.

    Examples
    --------
    >>> concat = Concatenate()
    >>> a = torch.randn(4, 10, 20)
    >>> b = torch.randn(4, 10, 30)
    >>> result = concat(a, b)
    >>> print(result.shape)
    torch.Size([4, 10, 50])
    """
    return torch.cat(inputs, dim=-1)

FeaturePartitioner

FeaturePartitioner(partitions: int, overlap: int)

Bases: Module

A layer that partitions the feature dimension into overlapping slices.

This layer is useful for dividing input features into structured regions while maintaining smooth transitions between partitions. Commonly used in parallel reservoir architectures where different reservoirs process different feature subspaces.

Behavior: - Splits the feature dimension into partitions groups - Each partition overlaps with its neighbors by overlap units - Applies circular wrapping: last overlap features wrap to start, and vice versa

Args: partitions: Number of partitions to divide the feature dimension into overlap: Overlap size (in feature units) for each partition

Input Shape: (..., features) — rank-agnostic on the feature (last) dimension. Both (batch, features) and (batch, sequence_length, features) are accepted; the leading dimensions are preserved. The 2-D rank lets the layer sit in the autoregressive forecast path, where the flattened engine feeds single-step slices.

Output: List of partitions tensors, each with the same leading dimensions as the input and a last dimension of partition_width = features // partitions + 2 * overlap

Raises: ValueError: At construction, if partitions < 1 or overlap < 0. ValueError: At forward, if features % partitions != 0 (unless partitions == 1). ValueError: At forward, if overlap >= features // partitions (invalid overlap size).

Warns: UserWarning: At construction, if partitions == 1 and overlap > 0, since the single-partition fast path returns the input unchanged and the overlap is silently ignored.

Example: >>> partitioner = FeaturePartitioner(partitions=2, overlap=1) >>> x = torch.arange(12).reshape(1, 1, 12).float() >>> outputs = partitioner(x) >>> len(outputs) 2 >>> outputs[0].shape torch.Size([1, 1, 8]) # 12//2 + 2*1 = 8

Configuration is validated eagerly so misconfiguration surfaces at construction rather than deep inside a forward pass (the layer lives in an eagerly-resolved symbolic graph). The divisibility and overlap-vs-width checks stay in :meth:forward, where the runtime feature count is known.

Args: partitions: Number of partitions (must be a positive integer) overlap: Overlap size between adjacent partitions (must be non-negative)

Raises: ValueError: If partitions < 1 or overlap < 0.

Warns: UserWarning: If partitions == 1 and overlap > 0 — the single-partition fast path returns the input unchanged, so the overlap has no effect.

Source code in src/resdag/layers/transforms/feature_partitioner.py

def __init__(self, partitions: int, overlap: int) -> None:
    """Initialize the FeaturePartitioner.

    Configuration is validated eagerly so misconfiguration surfaces at
    construction rather than deep inside a forward pass (the layer lives in
    an eagerly-resolved symbolic graph). The divisibility and
    overlap-vs-width checks stay in :meth:`forward`, where the runtime
    feature count is known.

    Args:
        partitions: Number of partitions (must be a positive integer)
        overlap: Overlap size between adjacent partitions (must be non-negative)

    Raises:
        ValueError: If ``partitions < 1`` or ``overlap < 0``.

    Warns:
        UserWarning: If ``partitions == 1`` and ``overlap > 0`` — the
            single-partition fast path returns the input unchanged, so the
            overlap has no effect.
    """
    super().__init__()
    if partitions < 1:
        raise ValueError(f"partitions must be a positive integer, got {partitions}")
    if overlap < 0:
        raise ValueError(f"overlap must be a non-negative integer, got {overlap}")
    if partitions == 1 and overlap > 0:
        warnings.warn(
            f"FeaturePartitioner: overlap={overlap} is ignored when partitions == 1; "
            f"the single-partition fast path returns the input unchanged.",
            stacklevel=2,
        )

    self.partitions = partitions
    self.overlap = overlap

forward

forward(input: Tensor) -> list[Tensor]

Split the feature dimension into overlapping partitions with circular wrapping.

Operates purely on the feature (last) dimension, so it is rank-agnostic: both 2-D (batch, features) and 3-D (batch, sequence_length, features) inputs are accepted (the latter is the usual sequence layout; the former is what the flattened single-step forecast engine feeds per step). The leading dimensions are preserved unchanged.

Args: input: Input tensor whose last dimension is the feature dimension, e.g. (batch, features) or (batch, sequence_length, features)

Returns: List of length self.partitions, each with the same leading dimensions as input and a last dimension of partition_width

Raises: ValueError: If feature dimension is not divisible by partitions ValueError: If overlap is too large relative to partition size

Source code in src/resdag/layers/transforms/feature_partitioner.py

def forward(self, input: torch.Tensor) -> list[torch.Tensor]:
    """Split the feature dimension into overlapping partitions with circular wrapping.

    Operates purely on the feature (last) dimension, so it is rank-agnostic:
    both 2-D ``(batch, features)`` and 3-D ``(batch, sequence_length, features)``
    inputs are accepted (the latter is the usual sequence layout; the former
    is what the flattened single-step forecast engine feeds per step). The
    leading dimensions are preserved unchanged.

    Args:
        input: Input tensor whose last dimension is the feature dimension,
            e.g. (batch, features) or (batch, sequence_length, features)

    Returns:
        List of length `self.partitions`, each with the same leading
        dimensions as `input` and a last dimension of `partition_width`

    Raises:
        ValueError: If feature dimension is not divisible by partitions
        ValueError: If overlap is too large relative to partition size
    """
    # If partitions == 1, just return the entire input as a single partition
    if self.partitions == 1:
        return [input]

    features = input.shape[-1]

    # Validate shape
    if features % self.partitions != 0:
        raise ValueError(
            f"Feature dimension ({features}) must be divisible by "
            f"number of partitions ({self.partitions})"
        )

    partition_base_width = features // self.partitions

    if self.overlap >= partition_base_width:
        raise ValueError(
            f"Overlap ({self.overlap}) must be smaller than the base partition "
            f"width ({partition_base_width})"
        )

    # Width of each partition including overlap
    partition_width = partition_base_width + 2 * self.overlap

    # Circular wrapping
    if self.overlap > 0:
        # Concatenate: [last overlap features | all features | first overlap features]
        wrapped_input = torch.cat(
            [input[..., -self.overlap :], input, input[..., : self.overlap]],
            dim=-1,
        )
    else:
        wrapped_input = input

    # Extract partitions
    partitions_out = []
    for i in range(self.partitions):
        start = i * partition_base_width
        end = start + partition_width
        partitions_out.append(wrapped_input[..., start:end])

    return partitions_out

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/layers/transforms/feature_partitioner.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    return f"partitions={self.partitions}, overlap={self.overlap}"

Power

Power(exponent: float, sign_preserving: bool = False)

Bases: Module

Exponentiate every feature to a fixed power.

Applies the chosen power element-wise along the last dimension. Used in power-augmented ESN architectures to enrich reservoir states before readout.

PARAMETER	DESCRIPTION
exponent	Power applied to each element of the input tensor. TYPE: float
sign_preserving	If True, every element uses the sign-preserving power sign(x) * abs(x) ** exponent. This keeps the forward and backward passes finite for negative bases under a non-integer exponent (e.g. tanh reservoir states in [-1, 1], which include negatives and zeros). If False (default), elements use plain :func:torch.pow, whose real-valued result is nan for a negative base raised to a non-integer exponent, and inf for a zero base raised to a negative exponent. TYPE: bool DEFAULT: False

Notes

With sign_preserving=False (the default), :func:torch.pow has two edge cases that silently corrupt readout inputs on common reservoir states:

• A negative base with a non-integer exponent produces nan in both the forward and backward pass (the real pow is undefined there). For example, Power(0.5) on [[-4.0, 4.0]] yields [[nan, 2.0]].
• A zero base with a negative exponent produces inf (division by zero). For example, Power(-1.0) on [[0.0, 2.0]] yields [[inf, 0.5]].

Tanh reservoir states live in [-1, 1] and routinely include negative and zero values, so a non-integer exponent on raw states is a plausible augmentation choice that silently emits nan/inf with no diagnostic. Use sign_preserving=True to apply sign(x) * abs(x) ** exponent, which stays finite for negative bases (the zero base with a negative exponent is still inf — abs(0) ** -1 diverges — and is unaffected by sign preservation). For integer exponents the two modes are identical for non-negative bases and differ only in sign for negative bases raised to an odd integer power. Every integer exponent — even ones map negatives to non-negative values, odd ones (including the power_augmented default exponent=3.0) preserve sign — is safe in either mode.

Examples:

>>> layer = Power(exponent=2.0)
>>> x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
>>> layer(x)
tensor([[ 1.,  4.,  9., 16.]])

Sign-preserving mode keeps negative bases finite under a fractional exponent, where plain :func:torch.pow would return nan:

>>> layer = Power(exponent=0.5, sign_preserving=True)
>>> x = torch.tensor([[-4.0, 4.0, -9.0, 9.0]])
>>> layer(x)
tensor([[-2.,  2., -3.,  3.]])

PARAMETER	DESCRIPTION
exponent	Value the input is raised to for every element. TYPE: float
sign_preserving	If True, apply sign(x) * abs(x) ** exponent so negative bases stay finite under a non-integer exponent. TYPE: bool DEFAULT: False

Source code in src/resdag/layers/transforms/power.py

def __init__(self, exponent: float, sign_preserving: bool = False) -> None:
    """Store the exponent and sign-preservation mode for the forward pass.

    Parameters
    ----------
    exponent : float
        Value the input is raised to for every element.
    sign_preserving : bool, default=False
        If ``True``, apply ``sign(x) * abs(x) ** exponent`` so negative
        bases stay finite under a non-integer exponent.
    """
    super().__init__()
    self.exponent = exponent
    self.sign_preserving = sign_preserving

forward

forward(input: Tensor) -> Tensor

Raise the input to self.exponent.

PARAMETER	DESCRIPTION
input	Tensor of shape (batch, ..., features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Same shape as input, with each element raised to exponent. When sign_preserving is set, sign(x) * abs(x) ** exponent is used, which is finite for negative bases under a non-integer exponent (where plain :func:torch.pow returns nan).

Source code in src/resdag/layers/transforms/power.py

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """Raise the input to ``self.exponent``.

    Parameters
    ----------
    input : torch.Tensor
        Tensor of shape ``(batch, ..., features)``.

    Returns
    -------
    torch.Tensor
        Same shape as *input*, with each element raised to ``exponent``.
        When ``sign_preserving`` is set, ``sign(x) * abs(x) ** exponent`` is
        used, which is finite for negative bases under a non-integer
        exponent (where plain :func:`torch.pow` returns ``nan``).
    """
    if self.sign_preserving:
        return input.sign() * input.abs().pow(self.exponent)
    return torch.pow(input, self.exponent)

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/layers/transforms/power.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    return f"exponent={self.exponent}, sign_preserving={self.sign_preserving}"

SelectiveDropout

SelectiveDropout(mask: Sequence[bool] | ndarray | Tensor)

Bases: Module

Layer that zeros out specific features based on a fixed mask.

This layer zeros out features at specified indices across all timesteps and batches. Unlike standard dropout which is stochastic, this layer uses a fixed mask provided at initialization. Useful for ablation studies and feature importance analysis.

Args: mask: Boolean mask where True indicates features to zero out. Can be array-like (list, numpy array) or torch.Tensor of shape (features,)

Input Shape: (batch, features) or (batch, timesteps, features) — the mask is applied on the feature (last) dimension, so both ranks are accepted. The autoregressive :meth:resdag.core.ESNModel.forecast engine drives the graph one timestep at a time, which is why per-step layers must accept the singleton-time slice as well as the full sequence.

Output Shape: Same shape as the input.

Example: >>> import numpy as np >>> mask = np.array([False, True, False, True]) # Drop indices 1 and 3 >>> layer = SelectiveDropout(mask) >>> x = torch.randn(2, 5, 4) >>> y = layer(x) >>> # Features at indices 1 and 3 are zeroed out

Args: mask: Boolean mask indicating which features to zero out (True = drop)

Raises: ValueError: If mask is not 1-dimensional

Source code in src/resdag/layers/transforms/selective_dropout.py

def __init__(self, mask: Sequence[bool] | np.ndarray | torch.Tensor) -> None:
    """Initialize the SelectiveDropout layer.

    Args:
        mask: Boolean mask indicating which features to zero out (True = drop)

    Raises:
        ValueError: If mask is not 1-dimensional
    """
    super().__init__()

    # Convert to torch tensor
    if isinstance(mask, torch.Tensor):
        mask_tensor = mask.bool()
    else:
        mask_tensor = torch.tensor(mask, dtype=torch.bool)

    # Validate shape
    if mask_tensor.ndim != 1:
        raise ValueError(f"Mask must be 1D, but got shape {mask_tensor.shape}")

    # Register as buffer (non-trainable, but part of state_dict)
    self.register_buffer("mask", mask_tensor)

forward

forward(input: Tensor) -> Tensor

Apply selective dropout using the stored mask.

Handles both 2-D (batch, features) and 3-D (batch, timesteps, features) inputs — the mask broadcasts over the feature (last) dimension, mirroring :class:~resdag.layers.readouts.ReadoutLayer. Accepting the 2-D rank lets the layer sit in the autoregressive forecast path, where the flattened engine feeds single-step slices.

Args: input: Input tensor of shape (batch, features) or (batch, timesteps, features)

Returns: Input tensor (same shape) with masked features set to zero

Raises: ValueError: If input is neither 2-D nor 3-D, or the mask size does not match the feature dimension

Source code in src/resdag/layers/transforms/selective_dropout.py

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """Apply selective dropout using the stored mask.

    Handles both 2-D ``(batch, features)`` and 3-D
    ``(batch, timesteps, features)`` inputs — the mask broadcasts over the
    feature (last) dimension, mirroring :class:`~resdag.layers.readouts.ReadoutLayer`.
    Accepting the 2-D rank lets the layer sit in the autoregressive forecast
    path, where the flattened engine feeds single-step slices.

    Args:
        input: Input tensor of shape (batch, features) or
            (batch, timesteps, features)

    Returns:
        Input tensor (same shape) with masked features set to zero

    Raises:
        ValueError: If input is neither 2-D nor 3-D, or the mask size does
            not match the feature dimension
    """
    if input.dim() not in (2, 3):
        raise ValueError(
            f"SelectiveDropout expects 2D (batch, features) or 3D "
            f"(batch, timesteps, features) input, got {input.dim()}D tensor "
            f"with shape {input.shape}"
        )

    mask = cast(torch.Tensor, self.mask)
    feature_dim = input.shape[-1]
    if mask.shape[0] != feature_dim:
        raise ValueError(
            f"Mask size ({mask.shape[0]}) does not match feature dimension ({feature_dim})"
        )

    # Apply mask: where mask is True, output 0; otherwise keep input
    return torch.where(mask, torch.zeros_like(input), input)

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/layers/transforms/selective_dropout.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    mask = cast(torch.Tensor, self.mask)
    num_dropped = mask.sum().item()
    total_features = mask.shape[0]
    return f"features={total_features}, dropped={num_dropped}"

SelectiveExponentiation

SelectiveExponentiation(index: int, exponent: float, sign_preserving: bool = False)

Bases: Module

Exponentiate even or odd feature indices based on parity.

If index is even, even positions in the last dimension are raised to exponent; if index is odd, odd positions are exponentiated. Other elements are unchanged. Used in Ott-style state-augmented ESNs.

With index=0, exponent=2 this is the T₁ / NLAT1 nonlinear algebraic transform — squaring every even-0-based feature — of Chattopadhyay, Hassanzadeh & Subramanian (2020) and Pathak et al. (2017), matching NLAT1 in ReservoirComputing.jl (which squares the odd entries of a 1-based feature-major state; those are the even-0-based features here). Its T₂/T₃ siblings are :class:~resdag.layers.transforms.NLAT2 and :class:~resdag.layers.transforms.NLAT3.

The transform is implemented with a :func:torch.where gate so that unselected positions never enter the pow node: they pass through with a gradient of 1 and are never poisoned by the pow backward (which is inf/nan at x = 0 for exponent < 1). The base fed to pow is also masked to a safe constant at unselected positions, which keeps unselected gradients finite even when those positions hold negative bases under a non-integer exponent.

The parity mask depends only on (feature_dim, device) — never on the input values or dtype — so it is computed once per (dim, device) pair and cached. This keeps the autoregressive forecast loop (which calls this layer once per step) off the per-step arange/parity/ones_like allocation path. The cache is a plain dict that stays out of state_dict and named_buffers; it is rebuilt lazily when a new dim or device is seen. For the common integer exponent=2 with sign_preserving=False the forward reduces to where(mask, x * x, x) (x * x is bit-identical to pow(x, 2) and finite everywhere, so the safe-base masking is unnecessary); the full safe-base path is retained for every other case.

PARAMETER	DESCRIPTION
index	Parity selector: even index targets even indices; odd index targets odd indices. TYPE: int
exponent	Power applied to the selected positions. TYPE: float
sign_preserving	If True, selected positions use the sign-preserving power sign(x) * abs(x) ** exponent. This keeps the forward and backward passes finite for negative selected bases under a non-integer exponent (e.g. tanh reservoir states in [-1, 1]). If False (default), selected positions use plain :func:torch.pow, whose real-valued result is nan for a negative base with a non-integer exponent. TYPE: bool DEFAULT: False

Notes

With sign_preserving=False and a non-integer exponent, selected positions holding a negative value produce nan in both the forward and backward pass — this is the mathematically correct behaviour of a real pow and matches :func:torch.pow. Unselected positions are always finite. For integer exponents (including the default ott_esn path, exponent=2.0) the two modes are identical for non-negative bases and differ only in sign for negative bases raised to an odd integer power.

Examples:

>>> layer = SelectiveExponentiation(index=2, exponent=2.0)
>>> x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
>>> layer(x)
tensor([[1., 2., 9., 4.]])

Sign-preserving mode keeps negative selected bases finite under a fractional exponent:

>>> layer = SelectiveExponentiation(index=0, exponent=0.5, sign_preserving=True)
>>> x = torch.tensor([[-4.0, 2.0, -9.0, 3.0]])
>>> layer(x)
tensor([[-2.,  2., -3.,  3.]])

PARAMETER	DESCRIPTION
index	Determines which feature indices are exponentiated (even vs odd). TYPE: int
exponent	Power applied to selected positions. TYPE: float
sign_preserving	If True, apply sign(x) * abs(x) ** exponent at selected positions so negative bases stay finite under a non-integer exponent. TYPE: bool DEFAULT: False

Source code in src/resdag/layers/transforms/selective_exponentiation.py

def __init__(self, index: int, exponent: float, sign_preserving: bool = False) -> None:
    """Store parity selector, exponent, and sign-preservation mode.

    Parameters
    ----------
    index : int
        Determines which feature indices are exponentiated (even vs odd).
    exponent : float
        Power applied to selected positions.
    sign_preserving : bool, default=False
        If ``True``, apply ``sign(x) * abs(x) ** exponent`` at selected
        positions so negative bases stay finite under a non-integer
        exponent.
    """
    super().__init__()
    self.index = index
    self.exponent = exponent
    self.sign_preserving = sign_preserving
    # Whether ``where(mask, x * x, x)`` reproduces the forward and gradient
    # exactly: only for a plain (non-sign-preserving) integer square. ``x * x``
    # is bit-identical to ``pow(x, 2)`` and finite for every real base, so the
    # safe-base masking is redundant there.
    self._is_square = (not sign_preserving) and float(exponent) == 2.0
    # Boolean parity masks cached per (feature_dim, device). Rebuilt lazily on
    # a new dim/device. Kept out of state_dict / named_buffers on purpose.
    self._mask_cache: dict[tuple[int, torch.device], torch.Tensor] = {}

forward

forward(input: Tensor) -> Tensor

Apply selective exponentiation along the last dimension.

PARAMETER	DESCRIPTION
input	Tensor of shape (batch, ..., features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Same shape as input; selected indices are raised to exponent, others are unchanged. Gradients at unselected positions are always finite; selected positions are finite when sign_preserving is set or the base/exponent combination is real-valued.

Source code in src/resdag/layers/transforms/selective_exponentiation.py

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """Apply selective exponentiation along the last dimension.

    Parameters
    ----------
    input : torch.Tensor
        Tensor of shape ``(batch, ..., features)``.

    Returns
    -------
    torch.Tensor
        Same shape as *input*; selected indices are raised to ``exponent``,
        others are unchanged. Gradients at unselected positions are always
        finite; selected positions are finite when ``sign_preserving`` is
        set or the base/exponent combination is real-valued.
    """
    mask = self._get_mask(input.shape[-1], input.device)

    # Fast path: a plain integer square is finite for every real base, so the
    # gate reduces to ``where(mask, x * x, x)``. This is bit-identical to the
    # safe-base path in both forward and gradient (``x * x == pow(x, 2)``) and
    # skips the ``ones_like`` allocation and the extra ``where``/``pow`` nodes.
    if self._is_square:
        return torch.where(mask, input * input, input)

    # General path. Feed a safe constant (1.0) into ``pow`` at unselected
    # positions so the pow node never evaluates a value that would
    # back-propagate as 0 * nan through torch.where (e.g. a negative
    # unselected base with a fractional exponent). pow(1, e) == 1 with a
    # finite gradient, and the where gate discards this branch for unselected
    # positions anyway.
    safe_base = torch.where(mask, input, torch.ones_like(input))

    if self.sign_preserving:
        powed = safe_base.sign() * safe_base.abs().pow(self.exponent)
    else:
        powed = safe_base.pow(self.exponent)

    return torch.where(mask, powed, input)

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/layers/transforms/selective_exponentiation.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    parity = "even" if self.index % 2 == 0 else "odd"
    return (
        f"index={self.index}, exponent={self.exponent}, "
        f"sign_preserving={self.sign_preserving}, applies_to={parity}_indices"
    )

Standardize

Standardize(num_features: int, mean: Tensor | None = None, std: Tensor | None = None, eps: float = 1e-08)

Bases: Module

Per-feature standardization (z-score) with a learnable-free inverse.

Centers and scales each feature of the last dimension using stored statistics, applying (x - mean) / std in :meth:forward and the exact inverse x * std + mean in :meth:inverse. The mean and std vectors are registered buffers of shape (num_features,), so they are part of the state_dict and move with .to()/dtype casts and the :meth:~resdag.core.ESNModel.save_full round trip.

Statistics may be supplied at construction or estimated from data with :meth:fit. Until either happens the layer initializes to mean = 0 and std = 1 (an identity transform).

A small eps is added to std before division to avoid blow-up on constant (zero-variance) features; the same eps-padded scale is used by :meth:inverse, so inverse(forward(x)) reconstructs x to within floating-point tolerance for every feature.

PARAMETER	DESCRIPTION
num_features	Size of the feature (last) dimension this layer standardizes. TYPE: int
mean	Initial per-feature mean of shape (num_features,). Defaults to zeros (no centering until :meth:fit is called). TYPE: Tensor DEFAULT: None
std	Initial per-feature standard deviation of shape (num_features,). Defaults to ones (no scaling until :meth:fit is called). TYPE: Tensor DEFAULT: None
eps	Numerical-stability floor added to std before dividing. Guards against division by zero on constant features. TYPE: float DEFAULT: 1e-8

ATTRIBUTE	DESCRIPTION
mean	Registered buffer holding the per-feature mean (num_features,). TYPE: Tensor
std	Registered buffer holding the per-feature standard deviation (num_features,). TYPE: Tensor

Examples:

Fit statistics from a batch, then standardize and invert:

>>> import torch
>>> from resdag.layers.transforms import Standardize
>>>
>>> layer = Standardize(num_features=3)
>>> x = torch.randn(8, 100, 3) * 5.0 + 2.0  # (batch, time, features)
>>> layer.fit(x)
Standardize(num_features=3, eps=1e-08)
>>> z = layer(x)  # ~zero mean, ~unit std per feature
>>> torch.allclose(layer.inverse(z), x, atol=1e-5)
True

Inside a composable pipeline (statistics travel with save_full/.to):

>>> import pytorch_symbolic as ps
>>> from resdag import ESNModel, ESNLayer, CGReadoutLayer, Standardize
>>>
>>> norm = Standardize(num_features=3)
>>> norm.fit(torch.randn(4, 200, 3))
Standardize(num_features=3, eps=1e-08)
>>> inp = ps.Input((200, 3))
>>> normed = norm(inp)
>>> reservoir = ESNLayer(100, feedback_size=3)(normed)
>>> readout = CGReadoutLayer(100, 3, name="output")(reservoir)
>>> model = ESNModel(inp, readout)

See Also

Power : Per-feature exponentiation. SelectiveExponentiation : Parity-selective exponentiation.

PARAMETER	DESCRIPTION
num_features	Size of the feature dimension to standardize. TYPE: int
mean	Initial per-feature mean of shape (num_features,) (default zeros). TYPE: Tensor DEFAULT: None
std	Initial per-feature standard deviation of shape (num_features,) (default ones). TYPE: Tensor DEFAULT: None
eps	Stability floor added to std before dividing. TYPE: float DEFAULT: 1e-8

RAISES	DESCRIPTION
ValueError	If num_features is not positive, or if a supplied mean/ std is not a 1-D tensor of length num_features.

Source code in src/resdag/layers/transforms/standardize.py

def __init__(
    self,
    num_features: int,
    mean: torch.Tensor | None = None,
    std: torch.Tensor | None = None,
    eps: float = 1e-8,
) -> None:
    """Register the per-feature mean/std buffers and the ``eps`` floor.

    Parameters
    ----------
    num_features : int
        Size of the feature dimension to standardize.
    mean : torch.Tensor, optional
        Initial per-feature mean of shape ``(num_features,)`` (default
        zeros).
    std : torch.Tensor, optional
        Initial per-feature standard deviation of shape
        ``(num_features,)`` (default ones).
    eps : float, default=1e-8
        Stability floor added to ``std`` before dividing.

    Raises
    ------
    ValueError
        If ``num_features`` is not positive, or if a supplied ``mean``/
        ``std`` is not a 1-D tensor of length ``num_features``.
    """
    super().__init__()

    if num_features <= 0:
        raise ValueError(f"num_features must be positive, got {num_features}")

    self.num_features = num_features
    self.eps = eps

    mean_buffer = self._validate_stat(mean, "mean", torch.zeros(num_features))
    std_buffer = self._validate_stat(std, "std", torch.ones(num_features))

    self.register_buffer("mean", mean_buffer)
    self.register_buffer("std", std_buffer)

fit

fit(input: Tensor) -> Standardize

Estimate per-feature mean and std from data and store them.

Statistics are reduced over every dimension except the last, so an input of shape (batch, time, features) (or any number of leading dims) yields per-feature vectors of length num_features. The population standard deviation (unbiased=False) is used so a single observation does not produce a nan. The buffers are updated in place and inherit the device/dtype of input.

PARAMETER	DESCRIPTION
input	Data of shape (..., num_features) whose last dimension matches num_features. TYPE: Tensor

RETURNS	DESCRIPTION
Standardize	self, so calls can be chained (e.g. layer.fit(x)(x)).

RAISES	DESCRIPTION
ValueError	If the last dimension of input does not equal num_features.

Source code in src/resdag/layers/transforms/standardize.py

@torch.no_grad()
def fit(self, input: torch.Tensor) -> "Standardize":
    """Estimate per-feature mean and std from data and store them.

    Statistics are reduced over every dimension except the last, so an
    input of shape ``(batch, time, features)`` (or any number of leading
    dims) yields per-feature vectors of length ``num_features``. The
    population standard deviation (``unbiased=False``) is used so a single
    observation does not produce a ``nan``. The buffers are updated in
    place and inherit the device/dtype of *input*.

    Parameters
    ----------
    input : torch.Tensor
        Data of shape ``(..., num_features)`` whose last dimension matches
        ``num_features``.

    Returns
    -------
    Standardize
        ``self``, so calls can be chained (e.g. ``layer.fit(x)(x)``).

    Raises
    ------
    ValueError
        If the last dimension of *input* does not equal ``num_features``.
    """
    if input.shape[-1] != self.num_features:
        raise ValueError(
            f"input feature dim ({input.shape[-1]}) does not match "
            f"num_features ({self.num_features})"
        )

    flat = input.reshape(-1, self.num_features).to(dtype=self.mean.dtype)
    self.mean.copy_(flat.mean(dim=0))
    self.std.copy_(flat.std(dim=0, unbiased=False))
    return self

forward

forward(input: Tensor) -> Tensor

Standardize the last dimension as (x - mean) / (std + eps).

PARAMETER	DESCRIPTION
input	Tensor of shape (..., num_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Same shape as input, centered and scaled per feature. Gradients flow through unchanged (the stored statistics are constants).

RAISES	DESCRIPTION
ValueError	If the last dimension of input does not equal num_features.

Source code in src/resdag/layers/transforms/standardize.py

def forward(self, input: torch.Tensor) -> torch.Tensor:
    """Standardize the last dimension as ``(x - mean) / (std + eps)``.

    Parameters
    ----------
    input : torch.Tensor
        Tensor of shape ``(..., num_features)``.

    Returns
    -------
    torch.Tensor
        Same shape as *input*, centered and scaled per feature. Gradients
        flow through unchanged (the stored statistics are constants).

    Raises
    ------
    ValueError
        If the last dimension of *input* does not equal ``num_features``.
    """
    if input.shape[-1] != self.num_features:
        raise ValueError(
            f"input feature dim ({input.shape[-1]}) does not match "
            f"num_features ({self.num_features})"
        )
    return (input - self.mean) / (self.std + self.eps)

inverse

inverse(input: Tensor) -> Tensor

Invert the transform as x * (std + eps) + mean.

The exact algebraic inverse of :meth:forward, using the same eps-padded scale so that inverse(forward(x)) == x up to floating-point error. Use it to map standardized model outputs (e.g. forecasts) back to the original data scale.

PARAMETER	DESCRIPTION
input	Standardized tensor of shape (..., num_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Same shape as input, mapped back to the original feature scale.

RAISES	DESCRIPTION
ValueError	If the last dimension of input does not equal num_features.

Source code in src/resdag/layers/transforms/standardize.py

def inverse(self, input: torch.Tensor) -> torch.Tensor:
    """Invert the transform as ``x * (std + eps) + mean``.

    The exact algebraic inverse of :meth:`forward`, using the same
    ``eps``-padded scale so that ``inverse(forward(x)) == x`` up to
    floating-point error. Use it to map standardized model outputs (e.g.
    forecasts) back to the original data scale.

    Parameters
    ----------
    input : torch.Tensor
        Standardized tensor of shape ``(..., num_features)``.

    Returns
    -------
    torch.Tensor
        Same shape as *input*, mapped back to the original feature scale.

    Raises
    ------
    ValueError
        If the last dimension of *input* does not equal ``num_features``.
    """
    if input.shape[-1] != self.num_features:
        raise ValueError(
            f"input feature dim ({input.shape[-1]}) does not match "
            f"num_features ({self.num_features})"
        )
    return input * (self.std + self.eps) + self.mean

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/layers/transforms/standardize.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    return f"num_features={self.num_features}, eps={self.eps}"