DynoNet Models

dynoNet linear-dynamical-operator model, with its private fused all-pole kernel.

The model lives in :mod:.core; :mod:.allpole_triton is the fused all-pole IIR denominator kernel it dispatches to directly.

LinearDynamicalOperator

LinearDynamicalOperator(in_channels: int, out_channels: int, nb: int = 8, na: int = 2, backend: str = 'scan')

Bases: Module

MIMO bank of learnable rational transfer functions G(q) = B(q) / A(q) (dynoNet G-block).

Each (output, input) channel pair owns an independent SISO filter with nb numerator taps b_0 .. b_{nb-1} and na monic-denominator coefficients a_1 .. a_na; output channels sum the filtered contributions of all inputs. The numerator is a grouped causal convolution; the denominator recurrence runs in state-space (companion) form through linear_recurrence, so the whole operator is exact and sequence-parallel.

Coefficients are unconstrained as in Forgione & Piga (2021, arXiv:2006.02250; full citation on DynoNet): b starts small and random, a starts at zero (all poles at the origin — a pure FIR filter), so the operator is stable at initialization but poles may leave the unit circle during training.

The internal pair flattening is input-major (index = j_in * out_channels + i_out), forced by conv1d group semantics; every reshape below relies on this ordering.

Parameters:

Name	Type	Description	Default
in_channels	int	number of input signals.	required
out_channels	int	number of output signals.	required
nb	int	number of numerator (FIR) taps per filter.	8
na	int	denominator order per filter; 0 gives a pure FIR operator.	2
backend	str	"scan" (parallel; on CUDA float32 the fused all-pole Triton kernel, honoring the process preference set via tsfast.models.set_backend/ use_backend, else the log-doubling matrix scan) or "eager" (sequential loop).	'scan'

Source code in tsfast/models/architectures/dynonet/core.py

def __init__(self, in_channels: int, out_channels: int, nb: int = 8, na: int = 2, backend: str = "scan"):
    super().__init__()
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.nb = nb
    self.na = na
    self.backend = backend
    self.b_coeff = nn.Parameter(torch.randn(out_channels, in_channels, nb) * 0.01)
    self.a_coeff = nn.Parameter(torch.zeros(out_channels, in_channels, na))

forward

forward(u: Tensor, state: dict | None = None, return_state: bool = False) -> torch.Tensor | tuple[torch.Tensor, dict]

Filter the input sequence through all channel pairs and sum over inputs.

Parameters:

Name	Type	Description	Default
u	Tensor	input sequence [batch, seq, in_channels].	required
state	dict | None	carried filter state {"u": FIR tail, "x": IIR states} from a previous chunk; zero initial conditions if None.	None
return_state	bool	if True, return (output, new_state).	False

Returns:

Type	Description
Tensor | tuple[Tensor, dict]	Output sequence [batch, seq, out_channels], optionally with the new state.

Source code in tsfast/models/architectures/dynonet/core.py

def forward(
    self, u: torch.Tensor, state: dict | None = None, return_state: bool = False
) -> torch.Tensor | tuple[torch.Tensor, dict]:
    """Filter the input sequence through all channel pairs and sum over inputs.

    Args:
        u: input sequence ``[batch, seq, in_channels]``.
        state: carried filter state ``{"u": FIR tail, "x": IIR states}`` from a previous
            chunk; zero initial conditions if None.
        return_state: if ``True``, return ``(output, new_state)``.

    Returns:
        Output sequence ``[batch, seq, out_channels]``, optionally with the new state.
    """
    B, L, _ = u.shape
    match state:
        case {"u": u_tail, "x": x0}:
            pass
        case None:
            u_tail = u.new_zeros(B, self.nb - 1, self.in_channels)
            x0 = u.new_zeros(B, self.n_pairs, self.na)
        case _:
            raise TypeError(f"expected state dict {{'u': tensor, 'x': tensor}}, got {type(state)}")

    u_buf = torch.cat((u_tail, u), dim=1)
    # conv1d computes cross-correlation, so the taps are flipped to realize b_0 u_t + ... + b_{nb-1} u_{t-nb+1};
    # the carried tail replaces the zero left-padding of a cold-started causal convolution.
    weight = self.b_coeff.permute(1, 0, 2).reshape(self.n_pairs, 1, self.nb).flip(-1)
    w = F.conv1d(u_buf.transpose(1, 2), weight, groups=self.in_channels)

    if self.na > 0:
        y_pairs = x_last = None
        match self.backend:
            case "scan":
                # the companion state is a shift register of past outputs, so the fused
                # kernel runs the scalar all-pole form y_t = w_t - sum_i a_i y_{t-i}
                a = self.a_coeff.permute(1, 0, 2).reshape(self.n_pairs, self.na)
                y_pairs = _fused_allpole(a, w, x0)
                if y_pairs is None:
                    x = linear_recurrence(self._companion(), F.pad(w.unsqueeze(-1), (0, self.na - 1)), x0)
            case "eager":
                x = _linear_recurrence_sequential(self._companion(), F.pad(w.unsqueeze(-1), (0, self.na - 1)), x0)
            case unknown:
                raise ValueError(f"unknown backend {unknown!r}, expected 'scan' or 'eager'")
        if y_pairs is None:
            y_pairs = x[..., 0]
            x_last = x[..., -1, :]
        else:
            # x_last[j] = y_{L-1-j}, drawing from x0 when the chunk is shorter than na
            x_last = (
                y_pairs[..., L - self.na :].flip(-1)
                if L >= self.na
                else torch.cat((y_pairs.flip(-1), x0[..., : self.na - L]), dim=-1)
            )
    else:
        y_pairs = w
        x_last = u.new_zeros(B, self.n_pairs, 0)

    y = y_pairs.view(B, self.in_channels, self.out_channels, L).sum(1).transpose(1, 2)
    if not return_state:
        return y
    new_state = {"u": u_buf[:, u_buf.shape[1] - (self.nb - 1) :], "x": x_last}
    return y, new_state

DynoNet

DynoNet(input_size: int, output_size: int, n_channels: int = 8, nb: int = 8, na: int = 2, hidden_size: int = 32, hidden_layers: int = 1, act: type[Module] = nn.Tanh, bypass: bool = True, backend: str = 'scan', return_state: bool = False)

Bases: Module

dynoNet: linear transfer-function blocks G interconnected with a static nonlinearity F.

Wiener-Hammerstein-like structure G1 -> F -> G2 with an optional parallel linear bypass path, the canonical architecture of Forgione & Piga (2021). F is a pointwise MLP (memoryless), so all dynamics live in the LinearDynamicalOperator blocks.

With return_state=True the model follows the stateful-model protocol (forward(u, state=...) -> (out, state)); the carried state holds each G-block's FIR tail and IIR states, so chunked rollouts are exactly equivalent to the full sequence and TbpttLearner works unchanged. Initial conditions are zero unless state is passed.

References

M. Forgione and D. Piga, "dynoNet: A neural network architecture for learning dynamical systems," International Journal of Adaptive Control and Signal Processing, 35(4):612-626, 2021. arXiv:2006.02250.

Parameters:

Name	Type	Description	Default
input_size	int	number of input signals.	required
output_size	int	number of output signals.	required
n_channels	int	signal width between the blocks.	8
nb	int	numerator taps per filter in every G-block.	8
na	int	denominator order per filter in every G-block.	2
hidden_size	int	hidden width of the static nonlinearity MLP.	32
hidden_layers	int	number of hidden layers of the static nonlinearity MLP.	1
act	type[Module]	activation class of the static nonlinearity MLP.	Tanh
bypass	bool	add a parallel linear path G_lin from input to output.	True
backend	str	execution backend of the G-blocks, see LinearDynamicalOperator.	'scan'
return_state	bool	if True, return (output, state) tuple.	False

Source code in tsfast/models/architectures/dynonet/core.py

def __init__(
    self,
    input_size: int,
    output_size: int,
    n_channels: int = 8,
    nb: int = 8,
    na: int = 2,
    hidden_size: int = 32,
    hidden_layers: int = 1,
    act: type[nn.Module] = nn.Tanh,
    bypass: bool = True,
    backend: str = "scan",
    return_state: bool = False,
):
    super().__init__()
    self.return_state = return_state
    self.g1 = LinearDynamicalOperator(input_size, n_channels, nb, na, backend)
    self.f = SeqLinear(n_channels, n_channels, hidden_size, hidden_layer=hidden_layers, act=act)
    self.g2 = LinearDynamicalOperator(n_channels, output_size, nb, na, backend)
    self.g_lin = LinearDynamicalOperator(input_size, output_size, nb, na, backend) if bypass else None

forward

forward(u: Tensor, state: dict | None = None) -> torch.Tensor | tuple[torch.Tensor, dict]

Run the block interconnection over the input sequence.

Parameters:

Name	Type	Description	Default
u	Tensor	input sequence [batch, seq, input_size].	required
state	dict | None	carried state {"g1": ..., "g2": ..., "lin": ...} from a previous chunk.	None

Returns:

Type	Description
Tensor | tuple[Tensor, dict]	Output sequence [batch, seq, output_size], or (sequence, state) when
Tensor | tuple[Tensor, dict]	return_state is set.

Source code in tsfast/models/architectures/dynonet/core.py

def forward(self, u: torch.Tensor, state: dict | None = None) -> torch.Tensor | tuple[torch.Tensor, dict]:
    """Run the block interconnection over the input sequence.

    Args:
        u: input sequence ``[batch, seq, input_size]``.
        state: carried state ``{"g1": ..., "g2": ..., "lin": ...}`` from a previous chunk.

    Returns:
        Output sequence ``[batch, seq, output_size]``, or ``(sequence, state)`` when
        ``return_state`` is set.
    """
    match state:
        case None:
            s1 = s2 = s_lin = None
        case dict():
            s1, s2, s_lin = state.get("g1"), state.get("g2"), state.get("lin")
        case _:
            raise TypeError(f"expected state dict, got {type(state)}")
    y1, s1 = self.g1(u, state=s1, return_state=True)
    y, s2 = self.g2(self.f(y1), state=s2, return_state=True)
    new_state = {"g1": s1, "g2": s2}
    if self.g_lin is not None:
        y_lin, s_lin = self.g_lin(u, state=s_lin, return_state=True)
        y = y + y_lin
        new_state["lin"] = s_lin
    if self.return_state:
        return y, new_state
    return y

linear_recurrence

linear_recurrence(A: Tensor, v: Tensor, x0: Tensor | None = None) -> torch.Tensor

Compute the linear recurrence x_t = A x_{t-1} + v_t with constant A via a log-doubling scan.

Because A is constant along the sequence, the recurrence is a prefix sum x_t = A^t x_0 + sum_k A^(t-k) v_k that parallelizes exactly: each doubling step extends the summation window by a factor of two using one batched matmul over the whole sequence, so the sequential depth is ceil(log2(L)) instead of L. Exact for any spectral radius of A. Gradients come from the analytic matrix adjoint (the reverse-time scan G_t = A^T G_{t+1} + g_t) rather than autograd replay through the doubling levels, so backward memory is O(L) instead of the O(L log L) the levels would retain. Real dtypes only. Runs as the tsfast::linear_recurrence custom op, so it composes with torch.compile.

Parameters:

Name	Type	Description	Default
A	Tensor	transition matrices [..., n, n], broadcast against the leading dims of v.	required
v	Tensor	input sequence [..., L, n].	required
x0	Tensor | None	initial state [..., n]; zeros if None.	None

Returns:

Type	Description
Tensor	States x_1 .. x_L as [..., L, n].

Source code in tsfast/models/architectures/dynonet/core.py

def linear_recurrence(A: torch.Tensor, v: torch.Tensor, x0: torch.Tensor | None = None) -> torch.Tensor:
    """Compute the linear recurrence ``x_t = A x_{t-1} + v_t`` with constant ``A`` via a log-doubling scan.

    Because ``A`` is constant along the sequence, the recurrence is a prefix sum
    ``x_t = A^t x_0 + sum_k A^(t-k) v_k`` that parallelizes exactly: each doubling step extends
    the summation window by a factor of two using one batched matmul over the whole sequence,
    so the sequential depth is ``ceil(log2(L))`` instead of ``L``. Exact for any spectral radius
    of ``A``. Gradients come from the analytic matrix adjoint (the reverse-time scan
    ``G_t = A^T G_{t+1} + g_t``) rather than autograd replay through the doubling levels, so
    backward memory is O(L) instead of the O(L log L) the levels would retain. Real dtypes only.
    Runs as the ``tsfast::linear_recurrence`` custom op, so it composes with ``torch.compile``.

    Args:
        A: transition matrices ``[..., n, n]``, broadcast against the leading dims of ``v``.
        v: input sequence ``[..., L, n]``.
        x0: initial state ``[..., n]``; zeros if None.

    Returns:
        States ``x_1 .. x_L`` as ``[..., L, n]``.
    """
    bshape = torch.broadcast_shapes(A.shape[:-2], v.shape[:-2], () if x0 is None else x0.shape[:-1])
    A_b = A.broadcast_to(bshape + A.shape[-2:])
    v_b = v.broadcast_to(bshape + v.shape[-2:])
    x0_b = None if x0 is None else x0.broadcast_to(bshape + x0.shape[-1:])
    return _linear_recurrence_op(A_b, v_b, x0_b)