REN / R2DN Models

Sequence models with certificates that hold at every value of their free parameters.

Two architectures from the same lineage, both feedback interconnections of a linear system with a nonlinearity, and both certified by construction rather than by projection, penalty or post-hoc verification:

• :mod:.core — the recurrent equilibrium network (:class:~.core.REN), whose nonlinearity is a layer of slope-restricted neurons with feedback onto itself. Contracting, Lipschitz and (Q,S,R)-dissipative variants.
• :mod:.r2dn — the robust recurrent deep network (:class:~.r2dn.R2DN), which drops that feedback and replaces the neurons with a 1-Lipschitz network (:mod:.lbdn). Same certificates, no equilibrium solve, and a certificate matrix whose size no longer grows with nonlinear capacity.

The static specs and the explicit-realization containers live in :mod:.common.

Both are laid out around one seam: the certificate construction runs once per forward and produces a bundle of plain tensors, and the sequential rollout reads nothing else. The (Q,S,R) variants and the fused rollout kernels attach on opposite sides of it without meeting — a kernel that gets ∂L/∂A right needs no opinion about ∂L/∂X, and autograd carries the rest.

The REN's rollout is sequential twice over, along the sequence and along the n_nl neurons, so the naive loop is dispatch-bound by two orders of magnitude. The fused backends collapse a whole rollout into one launch with a hand-derived BPTT backward (see MATH_REN.md):

• :mod:.backend_c: generic scalar-templated C++ (float and double), batch-parallel; the fp64 gradcheck vehicle and the fast CPU path.
• :mod:.backend_triton: persistent per-trajectory GPU kernel, float32, within the config caps; the fast CUDA training path.

Both run behind the tsfast::ren_rollout / tsfast::ren_rollout_train / tsfast::ren_rollout_bwd custom ops registered in :mod:.core and are selected through :class:~.core.REN's backend argument (or the process-wide preference from :func:tsfast.models.set_backend); each backend reports its own applicability via supports(spec, u, x0).

The R2DN gets the same treatment in :mod:.r2dn_backend_triton (MATH_R2DN.md, tsfast::r2dn_rollout*). Deleting the sweep makes its step cheap but leaves the launches untouched, so eager it is dispatch-bound just as the REN is; fused, its cost is flat in nonlinear capacity where the REN's still grows with n_nl, which is where the architecture's scalability claim finally shows.

RENParameterization

RENParameterization(spec: RENSpec, gamma: float = 1.0, qsr: tuple[Tensor, Tensor, Tensor] | None = None, eps: float = _EPS, polar: bool = True, init: str = 'long_memory')

Bases: Module

Free parameters of a REN and the direct construction of its explicit realization.

Every parameter here is unconstrained: the certificate is a property of the construction, not of where the optimizer happens to be. forward runs the whole construction — the XᵀX product, the partition of H, and the E⁻¹/Λ⁻¹ solves — and returns plain tensors.

Parameters:

Name	Type	Description	Default
spec	RENSpec	static architecture description.	required
gamma	float	certified incremental ℓ2 gain for variant="lipschitz". A runtime scalar: reassigning it re-derives the certificate on the next forward.	1.0
qsr	tuple[Tensor, Tensor, Tensor] | None	(Q, S, R) supply-rate matrices for variant="dissipative", shaped (ny,ny), (nu,ny), (nu,nu), with Q ≺ 0 and R - S Q⁻¹ Sᵀ ≻ 0.	None
eps	float	regularization floor on H.	_EPS
polar	bool	use the polar parameterization H = (p²/‖X‖²) XᵀX + εI, which decouples the scale of H from the direction of X.	True
init	str	"long_memory" builds X so the explicit A starts near the identity — random X yields fast-forgetting models that will not fit long horizons. "random" samples X directly.	'long_memory'

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

def __init__(
    self,
    spec: RENSpec,
    gamma: float = 1.0,
    qsr: tuple[Tensor, Tensor, Tensor] | None = None,
    eps: float = _EPS,
    polar: bool = True,
    init: str = "long_memory",
):
    super().__init__()
    if init not in ("long_memory", "random"):
        raise ValueError(f"unknown init {init!r}, expected 'long_memory' or 'random'")
    nx, nu, ny, nv = spec.n_state, spec.n_input, spec.n_output, spec.n_nl
    self.spec = spec
    self.gamma = gamma
    self.eps = eps
    self.polar = polar

    self.B2 = nn.Parameter(_lecun_normal_(torch.empty(nx, nu)))
    self.D12 = nn.Parameter(_lecun_normal_(torch.empty(nv, nu)))
    self.C2 = nn.Parameter(_lecun_normal_(torch.empty(ny, nx)))
    self.D21 = nn.Parameter(_lecun_normal_(torch.empty(ny, nv)))
    self.bx = nn.Parameter(torch.zeros(nx))
    self.bv = nn.Parameter(torch.zeros(nv))
    self.by = nn.Parameter(torch.zeros(ny))

    x = _long_memory_x(spec, eps) if init == "long_memory" else _lecun_normal_(torch.empty(spec.n_h, spec.n_h))
    self.X = nn.Parameter(x)
    self.Y1 = nn.Parameter(torch.eye(nx) if init == "long_memory" else _lecun_normal_(torch.empty(nx, nx)))
    self.p = nn.Parameter(x.pow(2).sum().add(eps).sqrt().reshape(1))

    if spec.variant == "contracting":
        self.D22 = nn.Parameter(torch.zeros(ny, nu))
    else:
        # D22 is no longer free: it is built from these through a nonsquare Cayley
        # transform so that ‖N‖ ≤ 1, which is what makes the supply rate's input
        # weight invertible. The values below give D22 = 0 at initialization.
        d = min(nu, ny)
        self.X3 = nn.Parameter(torch.eye(d))
        self.Y3 = nn.Parameter(torch.zeros(d, d))
        self.Z3 = nn.Parameter(torch.zeros(abs(ny - nu), d))

    if spec.variant == "dissipative":
        if qsr is None:
            raise ValueError("variant='dissipative' requires qsr=(Q, S, R)")
        q, s, r = (torch.as_tensor(m, dtype=torch.get_default_dtype()) for m in qsr)
        _check_qsr(q, s, r, nu, ny)
        self.register_buffer("Q", q)
        self.register_buffer("S", s)
        self.register_buffer("R", r)

forward

forward() -> ExplicitREN

Build the explicit realization from the current free parameters.

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

def forward(self) -> ExplicitREN:
    """Build the explicit realization from the current free parameters."""
    h, d22 = self._construct()
    spec = self.spec
    nx, nv, nu = spec.n_state, spec.n_nl, spec.n_input
    h11, h21, h22 = h[:nx, :nx], h[nx : nx + nv, :nx], h[nx : nx + nv, nx : nx + nv]
    h31, h32, h33 = h[nx + nv :, :nx], h[nx + nv :, nx : nx + nv], h[nx + nv :, nx + nv :]

    e = (h11 + h33 / spec.alpha**2 + self.Y1 - self.Y1.mH) / 2
    # One solve for every column block that E⁻¹ acts on; never form the inverse.
    a, b1, b2 = torch.linalg.solve(e, torch.cat((h31, h32, self.B2), dim=1)).split((nx, nv, nu), dim=1)
    lam_inv = (2.0 / torch.diagonal(h22)).unsqueeze(1)
    return ExplicitREN(
        A=a,
        B1=b1,
        B2=b2,
        C1=-lam_inv * h21,
        D11=-lam_inv * torch.tril(h22, -1),
        D12=lam_inv * self.D12,
        C2=self.C2,
        D21=self.D21,
        D22=d22,
        bx=self.bx,
        bv=self.bv,
        by=self.by,
    )

hmatrix

hmatrix() -> Tensor

The certificate matrix H, positive definite for any parameter values.

Exposed because it is the guarantee itself: the storage matrix P = H33 and the contraction LMI are read off its blocks.

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

def hmatrix(self) -> Tensor:
    """The certificate matrix ``H``, positive definite for any parameter values.

    Exposed because it is the guarantee itself: the storage matrix ``P = H33`` and the
    contraction LMI are read off its blocks.
    """
    return self._construct()[0]

qsr

qsr() -> tuple[Tensor, Tensor, Tensor]

Supply-rate matrices (Q, S, R) in effect.

variant="lipschitz" is the special case Q = -I/γ, S = 0, R = γI, so the certified gain is a runtime scalar rather than a stored matrix.

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

def qsr(self) -> tuple[Tensor, Tensor, Tensor]:
    """Supply-rate matrices ``(Q, S, R)`` in effect.

    ``variant="lipschitz"`` is the special case ``Q = -I/γ``, ``S = 0``, ``R = γI``, so
    the certified gain is a runtime scalar rather than a stored matrix.
    """
    ny, nu = self.spec.n_output, self.spec.n_input
    if self.spec.variant == "lipschitz":
        eye_y = torch.eye(ny, dtype=self.X.dtype, device=self.X.device)
        eye_u = torch.eye(nu, dtype=self.X.dtype, device=self.X.device)
        return -eye_y / self.gamma, self.X.new_zeros(nu, ny), self.gamma * eye_u
    # User-supplied matrices may sit right on the definiteness boundary; nudge them
    # inside it so the Cholesky factors the construction takes stay well conditioned.
    eye_q = torch.eye(ny, dtype=self.Q.dtype, device=self.Q.device)
    eye_r = torch.eye(nu, dtype=self.R.dtype, device=self.R.device)
    return self.Q - self.eps * eye_q, self.S, self.R + self.eps * eye_r

cache_key

cache_key() -> tuple

Identity of the current parameter values, for the inference-mode explicit cache.

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

def cache_key(self) -> tuple:
    """Identity of the current parameter values, for the inference-mode explicit cache."""
    return parameter_cache_key(self, self.gamma)

RENCore

RENCore(spec: RENSpec, **kwargs: Any)

Bases: Module

Explicit realization plus the sequential rollout over an input sequence.

Holds the free parameters (in :attr:parameterization) but evaluates only through :class:~.common.ExplicitREN tensors, which is the seam a fused kernel would attach to: the rollout has no opinion about how the matrices were certified.

Parameters:

Name	Type	Description	Default
spec	RENSpec	static architecture description.	required
**kwargs	Any	forwarded to :class:RENParameterization.	{}

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

def __init__(self, spec: RENSpec, **kwargs: Any):
    super().__init__()
    if spec.act not in _ACTS:
        raise ValueError(f"unknown activation {spec.act!r}, expected one of {sorted(_ACTS)}")
    if not 0.0 < spec.alpha <= 1.0:
        raise ValueError(f"alpha must lie in (0, 1], got {spec.alpha}")
    self.spec = spec
    self.parameterization = RENParameterization(spec, **kwargs)
    self._act = _ACTS[spec.act]
    self._cache: tuple[tuple, ExplicitREN] | None = None

explicit

explicit() -> ExplicitREN

The explicit realization, rebuilt on demand and cached while gradients are off.

The construction costs a few matrix products and two solves on (2nx+nv)-sized matrices — irrelevant next to an L-step rollout during training, but worth caching for repeated inference from fixed weights.

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

def explicit(self) -> ExplicitREN:
    """The explicit realization, rebuilt on demand and cached while gradients are off.

    The construction costs a few matrix products and two solves on ``(2nx+nv)``-sized
    matrices — irrelevant next to an ``L``-step rollout during training, but worth
    caching for repeated inference from fixed weights.
    """
    if torch.is_grad_enabled() or torch.compiler.is_compiling():
        return self.parameterization()
    key = self.parameterization.cache_key()
    if self._cache is None or self._cache[0] != key:
        self._cache = (key, self.parameterization())
    return self._cache[1]

rollout

rollout(e: ExplicitREN, u: Tensor, x0: Tensor) -> tuple[Tensor, Tensor]

Run the realization over u [B, L, n_input] from x0 [B, n_state].

Returns:

Type	Description
tuple[Tensor, Tensor]	(y, x_L) with y shaped [B, L, n_output].

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

def rollout(self, e: ExplicitREN, u: Tensor, x0: Tensor) -> tuple[Tensor, Tensor]:
    """Run the realization over ``u [B, L, n_input]`` from ``x0 [B, n_state]``.

    Returns:
        ``(y, x_L)`` with ``y`` shaped ``[B, L, n_output]``.
    """
    bv = u @ e.D12.mH + e.bv
    bx = u @ e.B2.mH + e.bx
    by = u @ e.D22.mH + e.by
    x = x0
    ys = []
    for t in range(u.shape[1]):
        w = equilibrium_sweep(x @ e.C1.mH + bv[:, t], e.D11, self._act)
        ys.append(x @ e.C2.mH + w @ e.D21.mH + by[:, t])
        x = x @ e.A.mH + w @ e.B1.mH + bx[:, t]
    return torch.stack(ys, dim=1), x

REN

REN(n_input: int, n_output: int, n_state: int = 8, n_nl: int = 32, variant: str = 'contracting', alpha: float = 1.0, gamma: float = 1.0, qsr: tuple[Tensor, Tensor, Tensor] | None = None, act: str = 'tanh', eps: float = _EPS, polar: bool = True, init: str = 'long_memory', backend: str = 'auto', return_state: bool = False)

Bases: Module

Recurrent equilibrium network: a sequence model that is contracting by construction.

Trained with plain SGD from any initialization, the model satisfies its certificate at every step of training, because the certificate is built into the map from free parameters to model matrices rather than enforced on top of it.

Three variants, differing only in how the certificate matrix is assembled:

• "contracting": two trajectories under the same input converge at rate alpha.
• "lipschitz": additionally ‖y(u) - y(ũ)‖ ≤ gamma ‖u - ũ‖ in truncated ℓ2 from a common initial state. gamma is a runtime scalar and may be reassigned.
• "dissipative": additionally satisfies the incremental IQC given by qsr.

What the Lipschitz certificate buys, precisely: an input perturbation of energy δ moves the output by at most gamma·δ, measured in ℓ2 over the horizon and starting from the same state. It says nothing about model-vs-plant error — a REN with gamma = 1 can be an arbitrarily bad model of a system, certified smooth and stable rather than correct.

The rollout is irreducibly sequential twice over — once along the sequence, once along the n_nl neurons of the equilibrium layer — so a naive Python loop is dispatch-bound by a wide margin. Several backends implement the identical recurrence:

• "eager": nested Python loop — the reference implementation, any device and dtype.
• "c": generated C++ rollout with a fused BPTT backward, batch-parallel via the ATen thread pool — float32 and float64 on CPU, and the gradcheck vehicle.
• "triton": persistent per-trajectory GPU kernel with a fused BPTT backward — float32 on CUDA, within the size caps its fits reports.
• "compiled": torch.compile over the unrolled loop. Only usable on short sequences — the graph has seq * n_nl nodes and compiles at roughly 0.2 s per node — so it is never selected implicitly.
• "auto": defers to the process-wide preference (tsfast.models.set_backend / use_backend); under an "auto" preference picks triton where it applies and eager elsewhere (select "c" explicitly to trade a one-time compilation for much faster CPU training). A "reference" preference forces the eager path everywhere.

All backends share the same parameters, so the backend can be switched at any time via the backend attribute. The fused backends run as registered torch.library custom ops with analytic-BPTT backward ops (MATH_REN.md), so they are loss-agnostic and compose with torch.compile. They consume the explicit realization as plain tensors and know nothing about the certificate, which is built in ordinary autograd above them.

Contraction is a prior about the plant, not free insurance, and it fits some systems badly. On benchmarks/gate_ren.py at matched parameter count it wins clearly on CascadedTanks (0.37 NRMSE against 0.60 for a GRU, and 0.10 against 0.23 with both under FranSys), ties on Silverbox, and loses by roughly a quarter on WH — where :class:~.r2dn.R2DN reaches 0.037 at the same budget and closes the gap, so that loss is the equilibrium layer's rather than contraction's. It loses badly on EMPS, by 4.2x under FranSys — that plant is friction-dominated, and stick-slip is exactly the behaviour a contraction certificate excludes, since trajectories in a stick phase do not converge. Do not reach for this model on stick-slip or hysteretic systems; the guarantee it offers is one those dynamics cannot satisfy in the first place.

gamma is not a free addition either. It costs nothing on WH (0.037, matching the contracting model) but a third of the accuracy on CascadedTanks (0.43 against 0.37), where the gain budget binds against dynamics that need the range. What it buys is a certificate that is tight on trained models — gamma_empirical/gamma_certified in 0.37-0.84 across the suite — rather than the orders-of-magnitude slack typical of post-hoc bounds on freely-trained networks.

Contraction makes the initial state self-correcting at rate alpha, so x0=None (zeros) plus n_skip is usually enough. But the forgetting time ≈ 1/(1-alpha) is the longest time constant the model can represent, so no alpha both forgets x0 quickly and represents an integrator. For integrating plants (position from velocity, tank level, thermal accumulation) use return_state=True and compose with :class:~tsfast.prediction.fransys.FranSys, which estimates x0 from an (u, y) window instead of asking the dynamics to forget it. Note that the Lipschitz bound is stated for a fixed initial state and does not survive that composition.

Parameters:

Name	Type	Description	Default
n_input	int	exogenous input dimension.	required
n_output	int	observed output dimension.	required
n_state	int	state dimension.	8
n_nl	int	neurons in the equilibrium layer; the model's nonlinear capacity.	32
variant	str	"contracting", "lipschitz" or "dissipative".	'contracting'
alpha	float	contraction rate in (0, 1]. 1.0 admits arbitrarily long memory.	1.0
gamma	float	certified incremental ℓ2 gain, for variant="lipschitz".	1.0
qsr	tuple[Tensor, Tensor, Tensor] | None	(Q, S, R) supply-rate matrices, required for variant="dissipative".	None
act	str	equilibrium-layer activation, one of tanh, relu, sigmoid; must be monotone and slope-restricted to [0, 1], which all three are.	'tanh'
eps	float	regularization floor on the certificate matrix.	_EPS
polar	bool	use the polar parameterization of H.	True
init	str	"long_memory" (default) or "random"; see :class:RENParameterization.	'long_memory'
backend	str	execution backend, see above.	'auto'
return_state	bool	if True, return (output, state) following the stateful-model protocol, so TbpttLearner state carrying and FranSys both work.	False

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

def __init__(
    self,
    n_input: int,
    n_output: int,
    n_state: int = 8,
    n_nl: int = 32,
    variant: str = "contracting",
    alpha: float = 1.0,
    gamma: float = 1.0,
    qsr: tuple[Tensor, Tensor, Tensor] | None = None,
    act: str = "tanh",
    eps: float = _EPS,
    polar: bool = True,
    init: str = "long_memory",
    backend: str = "auto",
    return_state: bool = False,
):
    super().__init__()
    if variant not in ("contracting", "lipschitz", "dissipative"):
        raise ValueError(f"unknown variant {variant!r}")
    spec = RENSpec(n_state, n_input, n_output, n_nl, variant, alpha, act)
    self.core = RENCore(spec, gamma=gamma, qsr=qsr, eps=eps, polar=polar, init=init)
    self.backend = backend
    self.return_state = return_state
    self._compiled_rollout = None

gamma property writable

gamma: float

Certified incremental ℓ2 gain; reassign to retune the certificate.

forward

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

Roll the certified dynamics over the input sequence.

Parameters:

Name	Type	Description	Default
u	Tensor	input sequence [batch, seq, n_input].	required
x0	Tensor | None	initial state [batch, n_state] (or [batch, 1, n_state]); zeros if None.	None
state	dict | None	carried state {"x": x_last} from a previous chunk; overrides x0.	None

Returns:

Type	Description
Tensor | tuple[Tensor, dict]	Output sequence [batch, seq, n_output] observing the states x_0 .. x_{L-1},
Tensor | tuple[Tensor, dict]	or (sequence, {"x": x_L}) when return_state is set.

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

def forward(self, u: Tensor, x0: Tensor | None = None, state: dict | None = None) -> Tensor | tuple[Tensor, dict]:
    """Roll the certified dynamics over the input sequence.

    Args:
        u: input sequence ``[batch, seq, n_input]``.
        x0: initial state ``[batch, n_state]`` (or ``[batch, 1, n_state]``); zeros if None.
        state: carried state ``{"x": x_last}`` from a previous chunk; overrides ``x0``.

    Returns:
        Output sequence ``[batch, seq, n_output]`` observing the states ``x_0 .. x_{L-1}``,
        or ``(sequence, {"x": x_L})`` when ``return_state`` is set.
    """
    match state:
        case {"x": x_carry}:
            x0 = x_carry
        case None:
            pass
        case _:
            raise TypeError(f"expected state dict {{'x': tensor}}, got {type(state)}")
    if x0 is None:
        x0 = u.new_zeros(u.shape[0], self.spec.n_state)
    elif x0.dim() == 3:
        x0 = x0.squeeze(1)
    match _rollout_mode(self.backend, self.spec, u, x0):
        case "eager":
            y, x_last = self.core.rollout(self.core.explicit(), u, x0)
        case "compiled":
            y, x_last = self._rollout_compiled(u, x0)
        case _:
            y, x_last = fused_rollout(self.spec, u, x0, self.core.explicit())
    if self.return_state:
        return y, {"x": x_last}
    return y

ExplicitSandwich dataclass

ExplicitSandwich(B: Tensor, bias: Tensor, A: Tensor | None = None, psi: Tensor | None = None)

The tensors a sandwich layer evaluates, once the Cayley transform has been taken.

A and psi are absent on the output layer, which is a plain norm-bounded linear map::

hidden: h ↦ √2 · A Ψ σ(√2 Ψ⁻¹ B h + bias)      with ‖[Aᵀ; Bᵀ]‖ an isometry
output: h ↦ B h + bias                          with ‖B‖ ≤ 1

SandwichLayer

SandwichLayer(n_in: int, n_out: int, act: str = 'relu', is_output: bool = False)

Bases: Module

A layer that is 1-Lipschitz at every value of its free parameters.

The weights come from a Cayley transform of one unconstrained matrix, which makes [Aᵀ; Bᵀ] an isometry, and the activation is sandwiched between Ψ and Ψ⁻¹. The bound then follows from the activation being slope-restricted to [0, 1] rather than from any norm product, so Ψ is free to rescale the units without spending gain budget — which is what keeps a deep stack expressive.

Parameters:

Name	Type	Description	Default
n_in	int	input features.	required
n_out	int	output features.	required
act	str	activation name; must be monotone and slope-restricted to [0, 1].	'relu'
is_output	bool	build the norm-bounded linear form instead, with no activation and no Ψ. Used for the last layer of an :class:LBDN.	False

Source code in tsfast/models/architectures/ren/lbdn.py

def __init__(self, n_in: int, n_out: int, act: str = "relu", is_output: bool = False):
    super().__init__()
    if act not in _ACTS:
        raise ValueError(f"unknown activation {act!r}, expected one of {sorted(_ACTS)}")
    self.n_in, self.n_out, self.is_output = n_in, n_out, is_output
    self._act = _ACTS[act]
    # column-stacked for the Cayley transform, so the fan-in is the row count
    self.XY = nn.Parameter(nn.init.normal_(torch.empty(n_in + n_out, n_out), std=(n_in + n_out) ** -0.5))
    self.a = nn.Parameter(self.XY.detach().pow(2).sum().add(_EPS).sqrt().reshape(1))
    self.b = nn.Parameter(torch.zeros(n_out))
    if not is_output:
        self.d = nn.Parameter(torch.zeros(n_out))

explicit

explicit() -> ExplicitSandwich

Take the Cayley transform of the free parameters.

Source code in tsfast/models/architectures/ren/lbdn.py

def explicit(self) -> ExplicitSandwich:
    """Take the Cayley transform of the free parameters."""
    a, b = cayley_blocks(self.XY, self.a, self.n_out)
    if self.is_output:
        return ExplicitSandwich(B=b, bias=self.b)
    return ExplicitSandwich(B=b, bias=self.b, A=a, psi=self.d.clamp(-_LOG_PSI_CLAMP, _LOG_PSI_CLAMP).exp())

forward

forward(h: Tensor, e: ExplicitSandwich | None = None) -> Tensor

Map h [..., n_in] to [..., n_out], from a prebuilt realization if given.

Source code in tsfast/models/architectures/ren/lbdn.py

def forward(self, h: Tensor, e: ExplicitSandwich | None = None) -> Tensor:
    """Map ``h [..., n_in]`` to ``[..., n_out]``, from a prebuilt realization if given."""
    return _sandwich_forward(self.explicit() if e is None else e, h, self._act)

LBDN

LBDN(n_input: int, n_output: int, hidden: tuple[int, ...] = (64, 64), act: str = 'relu', gamma: float = 1.0)

Bases: Module

Feedforward network with a certified Lipschitz bound of gamma.

A drop-in replacement for an MLP whose incremental gain is a design parameter: ‖f(u) - f(ũ)‖ ≤ gamma ‖u - ũ‖ holds for every value of the free parameters, so ordinary SGD cannot break it. The bound is on the map, not a bound on how well it fits.

gamma is a runtime scalar — reassign it and the next forward re-derives the network, since it only rescales the two ends of the stack.

Parameters:

Name	Type	Description	Default
n_input	int	input features.	required
n_output	int	output features.	required
hidden	tuple[int, ...]	width of each hidden layer; its length is the number of nonlinear layers.	(64, 64)
act	str	activation, one of tanh, relu, sigmoid.	'relu'
gamma	float	certified Lipschitz bound.	1.0

Source code in tsfast/models/architectures/ren/lbdn.py

def __init__(
    self,
    n_input: int,
    n_output: int,
    hidden: tuple[int, ...] = (64, 64),
    act: str = "relu",
    gamma: float = 1.0,
):
    super().__init__()
    if gamma <= 0:
        raise ValueError(f"gamma must be positive, got {gamma}")
    sizes = (n_input, *hidden, n_output)
    self.layers = nn.ModuleList(
        SandwichLayer(i, o, act, is_output=k == len(hidden)) for k, (i, o) in enumerate(zip(sizes[:-1], sizes[1:]))
    )
    self.gamma = gamma
    self._act = _ACTS[act]

explicit

explicit() -> tuple[ExplicitSandwich, ...]

Per-layer realizations, in evaluation order.

Identically shaped hidden layers — the usual case, and the only one a uniform width produces — take their Cayley transforms as a single batched call. The output layer has a different form and is always built on its own.

Source code in tsfast/models/architectures/ren/lbdn.py

def explicit(self) -> tuple[ExplicitSandwich, ...]:
    """Per-layer realizations, in evaluation order.

    Identically shaped hidden layers — the usual case, and the only one a uniform width
    produces — take their Cayley transforms as a single batched call. The output layer
    has a different form and is always built on its own.
    """
    hidden = self.layers[:-1]
    if len(hidden) > 1 and len({tuple(layer.XY.shape) for layer in hidden}) == 1:
        a, b = cayley_blocks(
            torch.stack([layer.XY for layer in hidden]),
            torch.stack([layer.a for layer in hidden]).unsqueeze(-1),
            hidden[0].n_out,
        )
        psi = torch.stack([layer.d for layer in hidden]).clamp(-_LOG_PSI_CLAMP, _LOG_PSI_CLAMP).exp()
        built = tuple(ExplicitSandwich(b[i], layer.b, a[i], psi[i]) for i, layer in enumerate(hidden))
    else:
        built = tuple(layer.explicit() for layer in hidden)
    return (*built, self.layers[-1].explicit())

forward

forward(h: Tensor, e: tuple[ExplicitSandwich, ...] | None = None) -> Tensor

Map h [..., n_input] to [..., n_output], from a prebuilt realization if given.

Source code in tsfast/models/architectures/ren/lbdn.py

def forward(self, h: Tensor, e: tuple[ExplicitSandwich, ...] | None = None) -> Tensor:
    """Map ``h [..., n_input]`` to ``[..., n_output]``, from a prebuilt realization if given."""
    return lbdn_forward(self.explicit() if e is None else e, h, self._act, self.gamma)

R2DNSpec dataclass

R2DNSpec(n_state: int, n_input: int, n_output: int, n_nl: int, hidden: tuple[int, ...], variant: str, alpha: float, act: str)

Static description of an R2DN.

The certified gain gamma is deliberately not a field: like the REN's it is a runtime scalar, so retuning it must not invalidate anything specialized on the spec.

Parameters:

Name	Type	Description	Default
n_state	int	state dimension nx.	required
n_input	int	exogenous input dimension nu.	required
n_output	int	observed output dimension ny.	required
n_nl	int	width of the interconnection, i.e. of both v and w.	required
hidden	tuple[int, ...]	hidden widths of the 1-Lipschitz network; its length is the network's depth.	required
variant	str	"contracting" or "lipschitz".	required
alpha	float	contraction rate ᾱ ∈ (0, 1].	required
act	str	activation name; must be monotone and slope-restricted to [0, 1].	required

n_h property

n_h: int

Side length of the certificate matrix H: 2*n_state.

Independent of n_nl and of the network's depth, which is the scalability claim — the REN's H is (2*n_state + n_nl)² and grows with nonlinear capacity.

ExplicitR2DN dataclass

ExplicitR2DN(A: Tensor, B1: Tensor, B2: Tensor, C1: Tensor, C2: Tensor, D12: Tensor, D21: Tensor, D22: Tensor, bx: Tensor, bv: Tensor, by: Tensor, net: tuple[ExplicitSandwich, ...])

The realization the rollout actually runs, as plain tensors.

Per sample, with φ the 1-Lipschitz network held in :attr:net::

w   = φ(x C1ᵀ + u D12ᵀ + bv)
x⁺  = x Aᵀ + w B1ᵀ + u B2ᵀ + bx
y   = x C2ᵀ + w D21ᵀ + u D22ᵀ + by

y observes the state before the update, so a rollout of length L from x0 returns y_0 .. y_{L-1} and carries x_L.

tensors property

tensors: list[Tensor]

The linear part, flattened in field order; excludes :attr:net.

R2DNParameterization

R2DNParameterization(spec: R2DNSpec, gamma: float = 1.0, eps: float = _EPS, polar: bool = True, init: str = 'long_memory')

Bases: Module

Free parameters of an R2DN and the direct construction of its explicit realization.

Every parameter is unconstrained. The construction places the certificate's given terms into H = XᵀX + εI + (given) so that the dissipation LMI it has to satisfy reduces to XᵀX + εI ≻ 0 — true for any X — and then reads the realization off the blocks of H. The free B1, B2 are the implicit input maps E B1, E B2, since the LMI constrains those rather than the explicit ones.

Parameters:

Name	Type	Description	Default
spec	R2DNSpec	static architecture description.	required
gamma	float	certified incremental ℓ2 gain for variant="lipschitz". A runtime scalar: reassigning it re-derives the certificate on the next forward.	1.0
eps	float	regularization floor on H.	_EPS
polar	bool	use the polar parameterization H = (p²/‖X‖²) XᵀX + εI + ..., which decouples the scale of H from the direction of X.	True
init	str	"long_memory" builds X so the explicit A starts near the identity — random X yields fast-forgetting models that will not fit long horizons. "random" samples X directly.	'long_memory'

Source code in tsfast/models/architectures/ren/r2dn.py

def __init__(
    self,
    spec: R2DNSpec,
    gamma: float = 1.0,
    eps: float = _EPS,
    polar: bool = True,
    init: str = "long_memory",
):
    super().__init__()
    if init not in ("long_memory", "random"):
        raise ValueError(f"unknown init {init!r}, expected 'long_memory' or 'random'")
    nx, nu, ny, nv = spec.n_state, spec.n_input, spec.n_output, spec.n_nl
    self.spec = spec
    self.gamma = gamma
    self.eps = eps
    self.polar = polar

    # gamma = 1 is not a default here but the certificate's premise: the LMI below is
    # built around ‖Δw‖ ≤ ‖Δv‖, and a network of any other gain would break it
    self.net = LBDN(nv, nv, spec.hidden, act=spec.act, gamma=1.0)

    self.B2 = nn.Parameter(_lecun_normal_(torch.empty(nx, nu)))
    self.C2 = nn.Parameter(_lecun_normal_(torch.empty(ny, nx)))
    self.bx = nn.Parameter(torch.zeros(nx))
    self.bv = nn.Parameter(torch.zeros(nv))
    self.by = nn.Parameter(torch.zeros(ny))

    long_memory = init == "long_memory"
    x = _long_memory_x(spec, eps) if long_memory else _lecun_normal_(torch.empty(spec.n_h, spec.n_h))
    self.X = nn.Parameter(x)
    self.p = nn.Parameter(x.pow(2).sum().add(eps).sqrt().reshape(1))
    # the long-memory target realization assumes a dead nonlinear coupling and Y = E = I
    self.Y = nn.Parameter(torch.eye(nx) if long_memory else _lecun_normal_(torch.empty(nx, nx)))
    self.B1 = nn.Parameter(torch.zeros(nx, nv) if long_memory else _lecun_normal_(torch.empty(nx, nv)))
    self.C1 = nn.Parameter(torch.zeros(nv, nx) if long_memory else _lecun_normal_(torch.empty(nv, nx)))

    if spec.variant == "contracting":
        self.D12 = nn.Parameter(_lecun_normal_(torch.empty(nv, nu)))
        self.D21 = nn.Parameter(_lecun_normal_(torch.empty(ny, nv)))
        self.D22 = nn.Parameter(torch.zeros(ny, nu))
    else:
        # D12 and D21 stop being free: the supply rate's input weight R is only positive
        # definite while both stay below √gamma in norm, so each is built from a
        # Cayley transform of these. D22 is dropped entirely (see the class docstring).
        self._cayley_params("12", nv, nu)
        self._cayley_params("21", ny, nv)

forward

forward() -> ExplicitR2DN

Build the explicit realization from the current free parameters.

Source code in tsfast/models/architectures/ren/r2dn.py

def forward(self) -> ExplicitR2DN:
    """Build the explicit realization from the current free parameters."""
    h, d12, d21, d22 = self._construct()
    nx, nv, nu = self.spec.n_state, self.spec.n_nl, self.spec.n_input
    h11, h21, h22 = h[:nx, :nx], h[nx:, :nx], h[nx:, nx:]

    e = (h11 + h22 / self.spec.alpha**2 + self.Y - self.Y.mH) / 2
    # One solve for every column block that E⁻¹ acts on; never form the inverse.
    a, b1, b2 = torch.linalg.solve(e, torch.cat((h21, self.B1, self.B2), dim=1)).split((nx, nv, nu), dim=1)
    return ExplicitR2DN(
        A=a,
        B1=b1,
        B2=b2,
        C1=self.C1,
        C2=self.C2,
        D12=d12,
        D21=d21,
        D22=d22,
        bx=self.bx,
        bv=self.bv,
        by=self.by,
        net=self.net.explicit(),
    )

hmatrix

hmatrix() -> Tensor

The certificate matrix H, positive definite for any parameter values.

Exposed because it is the guarantee itself: the storage matrix P = H22 and the dissipation LMI are read off its blocks.

Source code in tsfast/models/architectures/ren/r2dn.py

def hmatrix(self) -> Tensor:
    """The certificate matrix ``H``, positive definite for any parameter values.

    Exposed because it is the guarantee itself: the storage matrix ``P = H22`` and the
    dissipation LMI are read off its blocks.
    """
    return self._construct()[0]

cache_key

cache_key() -> tuple

Identity of the current parameter values, for the inference-mode explicit cache.

Source code in tsfast/models/architectures/ren/r2dn.py

def cache_key(self) -> tuple:
    """Identity of the current parameter values, for the inference-mode explicit cache."""
    return parameter_cache_key(self, self.gamma)

R2DNCore

R2DNCore(spec: R2DNSpec, **kwargs: Any)

Bases: Module

Explicit realization plus the sequential rollout over an input sequence.

Holds the free parameters (in :attr:parameterization) but evaluates only through :class:ExplicitR2DN tensors: the rollout has no opinion about how the matrices were certified, and the 1-Lipschitz network's Cayley transforms are taken once per rollout rather than once per timestep.

Parameters:

Name	Type	Description	Default
spec	R2DNSpec	static architecture description.	required
**kwargs	Any	forwarded to :class:R2DNParameterization.	{}

Source code in tsfast/models/architectures/ren/r2dn.py

def __init__(self, spec: R2DNSpec, **kwargs: Any):
    super().__init__()
    if spec.act not in _ACTS:
        raise ValueError(f"unknown activation {spec.act!r}, expected one of {sorted(_ACTS)}")
    if not 0.0 < spec.alpha <= 1.0:
        raise ValueError(f"alpha must lie in (0, 1], got {spec.alpha}")
    self.spec = spec
    self.parameterization = R2DNParameterization(spec, **kwargs)
    self._act = _ACTS[spec.act]
    self._cache: tuple[tuple, ExplicitR2DN] | None = None

explicit

explicit() -> ExplicitR2DN

The explicit realization, rebuilt on demand and cached while gradients are off.

The construction costs a few matrix products, one solve on a 2*n_state matrix and one Cayley transform per network layer — irrelevant next to an L-step rollout during training, but worth caching for repeated inference from fixed weights.

Source code in tsfast/models/architectures/ren/r2dn.py

def explicit(self) -> ExplicitR2DN:
    """The explicit realization, rebuilt on demand and cached while gradients are off.

    The construction costs a few matrix products, one solve on a ``2*n_state`` matrix and
    one Cayley transform per network layer — irrelevant next to an ``L``-step rollout
    during training, but worth caching for repeated inference from fixed weights.
    """
    if torch.is_grad_enabled() or torch.compiler.is_compiling():
        return self.parameterization()
    key = self.parameterization.cache_key()
    if self._cache is None or self._cache[0] != key:
        self._cache = (key, self.parameterization())
    return self._cache[1]

rollout

rollout(e: ExplicitR2DN, u: Tensor, x0: Tensor) -> tuple[Tensor, Tensor]

Run the realization over u [B, L, n_input] from x0 [B, n_state].

Returns:

Type	Description
tuple[Tensor, Tensor]	(y, x_L) with y shaped [B, L, n_output].

Source code in tsfast/models/architectures/ren/r2dn.py

def rollout(self, e: ExplicitR2DN, u: Tensor, x0: Tensor) -> tuple[Tensor, Tensor]:
    """Run the realization over ``u [B, L, n_input]`` from ``x0 [B, n_state]``.

    Returns:
        ``(y, x_L)`` with ``y`` shaped ``[B, L, n_output]``.
    """
    bv = u @ e.D12.mH + e.bv
    bx = u @ e.B2.mH + e.bx
    by = u @ e.D22.mH + e.by
    x = x0
    ys = []
    for t in range(u.shape[1]):
        w = lbdn_forward(e.net, x @ e.C1.mH + bv[:, t], self._act)
        ys.append(x @ e.C2.mH + w @ e.D21.mH + by[:, t])
        x = x @ e.A.mH + w @ e.B1.mH + bx[:, t]
    return torch.stack(ys, dim=1), x

R2DN

R2DN(n_input: int, n_output: int, n_state: int = 8, n_nl: int = 32, depth: int = 2, variant: str = 'contracting', alpha: float = 1.0, gamma: float = 1.0, act: str = 'relu', eps: float = _EPS, polar: bool = True, init: str = 'long_memory', backend: str = 'auto', return_state: bool = False)

Bases: Module

Robust recurrent deep network: contracting by construction, with a deep nonlinearity.

Trained with plain SGD from any initialization, the model satisfies its certificate at every step of training, because the certificate is built into the map from free parameters to model matrices rather than enforced on top of it.

Two variants, differing only in how the certificate matrix is assembled:

• "contracting": two trajectories under the same input converge at rate alpha.
• "lipschitz": additionally ‖y(u) - y(ũ)‖ ≤ gamma ‖u - ũ‖ in truncated ℓ2 from a common initial state. gamma is a runtime scalar and may be reassigned. The paper's parameterization of this case covers the subset with D22 = 0, so there is no direct feedthrough from input to output; the nonlinearity still provides a static path through D12 and D21.

What the Lipschitz certificate buys, precisely: an input perturbation of energy δ moves the output by at most gamma·δ, measured in ℓ2 over the horizon and starting from the same state. It says nothing about model-vs-plant error — an R2DN with gamma = 1 can be an arbitrarily bad model of a system, certified smooth and stable rather than correct.

Against :class:~.core.REN, which certifies the same properties: nonlinear capacity here is depth in a 1-Lipschitz network instead of width in an equilibrium layer, so a step costs depth small GEMMs rather than a sequential sweep over n_nl neurons, and the certificate matrix stays 2·n_state square however large the nonlinearity grows. The REN's parameterization is the more general one — its equilibrium layer contains multi-layer networks as special cases — but the sweep is what makes it expensive.

Both models are dispatch-bound before they are fused, and both have a kernel that fixes that, so the comparison is between the fused paths. There the scalability claim holds: this rollout's cost is flat in nonlinear capacity while the REN's grows with n_nl, so at matched parameter count the two cross over around n_nl ≈ 24 and the R2DN is several times faster beyond it — and slower below, where the REN's sweep is only a few neurons long. benchmarks/benchmark_r2dn.py times both at matched parameter count.

The rollout is sequential along the sequence, so a Python loop is dispatch-bound at short n_state:

• "eager": the loop — any device and dtype.
• "triton": persistent per-trajectory GPU kernel with a fused BPTT backward (MATH_R2DN.md) — float32 on CUDA, within the size caps its fits reports.
• "compiled": torch.compile over the unrolled loop. Only usable on short sequences (the graph holds seq copies of the network), so it is never selected implicitly.
• "auto": defers to the process-wide preference (tsfast.models.set_backend / use_backend); under an "auto" preference picks triton where it applies and eager elsewhere. A "reference" preference forces the eager path everywhere.

All backends share the same parameters, so the backend can be switched at any time via the backend attribute. There is no CPU kernel: a "c" request selects eager.

Contraction is a prior about the plant, not free insurance, and it fits some systems badly: on friction-dominated plants the stick-slip phases are exactly what a contraction certificate excludes, since trajectories in a stick phase do not converge. On benchmarks/gate_ren.py that shows up on EMPS, where the best certified model trails a GRU by 2.3x under FranSys — better than the REN manages there, but the same verdict. Do not reach for either model on stick-slip or hysteretic plants.

Against the REN at matched parameter count the two are at parity overall, which is what the paper reports: this model wins on WH (0.037 NRMSE against 0.045, closing the one benchmark the REN loses to a GRU) and on EMPS under FranSys (0.058 against 0.106), ties on Silverbox and on CascadedTanks under FranSys (0.097 both), and loses on both integrating plants standalone (0.43 against 0.37, 0.36 against 0.28). Prefer it when the nonlinearity has to be large, where its cost is flat in capacity and the REN's is not; prefer the REN at small n_nl.

gamma is not a free addition. It costs nothing on WH (0.038 against 0.037 contracting) but a great deal on the integrating plants (0.61 against 0.43 on CascadedTanks), where the budget binds against dynamics that need the range. Prescribe it when something downstream consumes the bound, not by default. What it does buy is honest: on trained models the certificate is tight, gamma_empirical/gamma_certified landing in 0.48-0.83 across the suite — the same range the REN reaches, and nothing like the orders-of-magnitude slack typical of post-hoc bounds on freely-trained networks.

Contraction makes the initial state self-correcting at rate alpha, so x0=None (zeros) plus n_skip is usually enough. But the forgetting time ≈ 1/(1-alpha) is the longest time constant the model can represent, so no alpha both forgets x0 quickly and represents an integrator. For integrating plants (position from velocity, tank level, thermal accumulation) use return_state=True and compose with :class:~tsfast.prediction.fransys.FranSys, which estimates x0 from an (u, y) window instead of asking the dynamics to forget it. Note that the Lipschitz bound is stated for a fixed initial state and does not survive that composition.

Parameters:

Name	Type	Description	Default
n_input	int	exogenous input dimension.	required
n_output	int	observed output dimension.	required
n_state	int	state dimension.	8
n_nl	int	width of the interconnection with the nonlinearity, and of its hidden layers.	32
depth	int	number of nonlinear layers in the 1-Lipschitz network; at least 1.	2
variant	str	"contracting" or "lipschitz".	'contracting'
alpha	float	contraction rate in (0, 1]. 1.0 admits arbitrarily long memory.	1.0
gamma	float	certified incremental ℓ2 gain, for variant="lipschitz".	1.0
act	str	activation of the 1-Lipschitz network, one of tanh, relu, sigmoid; must be monotone and slope-restricted to [0, 1], which all three are.	'relu'
eps	float	regularization floor on the certificate matrix.	_EPS
polar	bool	use the polar parameterization of H.	True
init	str	"long_memory" (default) or "random"; see :class:R2DNParameterization.	'long_memory'
backend	str	execution backend, see above.	'auto'
return_state	bool	if True, return (output, state) following the stateful-model protocol, so TbpttLearner state carrying and FranSys both work.	False

Source code in tsfast/models/architectures/ren/r2dn.py

def __init__(
    self,
    n_input: int,
    n_output: int,
    n_state: int = 8,
    n_nl: int = 32,
    depth: int = 2,
    variant: str = "contracting",
    alpha: float = 1.0,
    gamma: float = 1.0,
    act: str = "relu",
    eps: float = _EPS,
    polar: bool = True,
    init: str = "long_memory",
    backend: str = "auto",
    return_state: bool = False,
):
    super().__init__()
    if variant not in ("contracting", "lipschitz"):
        raise ValueError(f"unknown variant {variant!r}, expected 'contracting' or 'lipschitz'")
    if depth < 1:
        raise ValueError(f"depth must be at least 1, got {depth}")
    spec = R2DNSpec(n_state, n_input, n_output, n_nl, (n_nl,) * depth, variant, alpha, act)
    self.core = R2DNCore(spec, gamma=gamma, eps=eps, polar=polar, init=init)
    self.backend = backend
    self.return_state = return_state
    self._compiled_rollout = None

gamma property writable

gamma: float

Certified incremental ℓ2 gain; reassign to retune the certificate.

forward

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

Roll the certified dynamics over the input sequence.

Parameters:

Name	Type	Description	Default
u	Tensor	input sequence [batch, seq, n_input].	required
x0	Tensor | None	initial state [batch, n_state] (or [batch, 1, n_state]); zeros if None.	None
state	dict | None	carried state {"x": x_last} from a previous chunk; overrides x0.	None

Returns:

Type	Description
Tensor | tuple[Tensor, dict]	Output sequence [batch, seq, n_output] observing the states x_0 .. x_{L-1},
Tensor | tuple[Tensor, dict]	or (sequence, {"x": x_L}) when return_state is set.

Source code in tsfast/models/architectures/ren/r2dn.py

def forward(self, u: Tensor, x0: Tensor | None = None, state: dict | None = None) -> Tensor | tuple[Tensor, dict]:
    """Roll the certified dynamics over the input sequence.

    Args:
        u: input sequence ``[batch, seq, n_input]``.
        x0: initial state ``[batch, n_state]`` (or ``[batch, 1, n_state]``); zeros if None.
        state: carried state ``{"x": x_last}`` from a previous chunk; overrides ``x0``.

    Returns:
        Output sequence ``[batch, seq, n_output]`` observing the states ``x_0 .. x_{L-1}``,
        or ``(sequence, {"x": x_L})`` when ``return_state`` is set.
    """
    match state:
        case {"x": x_carry}:
            x0 = x_carry
        case None:
            pass
        case _:
            raise TypeError(f"expected state dict {{'x': tensor}}, got {type(state)}")
    if x0 is None:
        x0 = u.new_zeros(u.shape[0], self.spec.n_state)
    elif x0.dim() == 3:
        x0 = x0.squeeze(1)
    match _rollout_mode(self.backend, self.spec, u, x0):
        case "eager":
            y, x_last = self.core.rollout(self.core.explicit(), u, x0)
        case "compiled":
            y, x_last = self._rollout_compiled(u, x0)
        case _:
            y, x_last = fused_rollout(self.spec, u, x0, self.core.explicit())
    if self.return_state:
        return y, {"x": x_last}
    return y

equilibrium_sweep

equilibrium_sweep(b: Tensor, d11: Tensor, act: Callable[[Tensor], Tensor]) -> Tensor

Solve w = act(w D11ᵀ + b) by forward substitution over the neurons.

D11 is strictly lower triangular by construction, so neuron i depends only on 0 .. i-1 and the equilibrium resolves exactly in one sweep. Each step is a rank-1 update of the pending pre-activations, which is why the loop costs two tensor ops per neuron rather than a growing matmul.

Parameters:

Name	Type	Description	Default
b	Tensor	input-and-state part of the pre-activation [..., n_nl].	required
d11	Tensor	strictly lower triangular feedback [n_nl, n_nl].	required
act	Callable[[Tensor], Tensor]	the activation, monotone and slope-restricted to [0, 1].	required

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

def equilibrium_sweep(b: Tensor, d11: Tensor, act: Callable[[Tensor], Tensor]) -> Tensor:
    """Solve ``w = act(w D11ᵀ + b)`` by forward substitution over the neurons.

    ``D11`` is strictly lower triangular by construction, so neuron ``i`` depends only on
    ``0 .. i-1`` and the equilibrium resolves exactly in one sweep. Each step is a rank-1
    update of the pending pre-activations, which is why the loop costs two tensor ops per
    neuron rather than a growing matmul.

    Args:
        b: input-and-state part of the pre-activation ``[..., n_nl]``.
        d11: strictly lower triangular feedback ``[n_nl, n_nl]``.
        act: the activation, monotone and slope-restricted to ``[0, 1]``.
    """
    acc = b
    ws = []
    for i in range(d11.shape[0]):
        w_i = act(acc[..., i])
        ws.append(w_i)
        if i + 1 < d11.shape[0]:
            acc = acc + w_i.unsqueeze(-1) * d11[:, i]
    return torch.stack(ws, dim=-1)

fused_rollout

fused_rollout(spec: RENSpec, u: Tensor, x0: Tensor, e: ExplicitREN) -> tuple[Tensor, Tensor]

Run the rollout through the fused-kernel custom ops (autograd-capable).

Picks the training op (which stores the states and equilibrium activations for the analytic BPTT backward) when gradients are live, else the inference op, which keeps no intermediates.

Returns:

Type	Description
tuple[Tensor, Tensor]	(y, x_L) with y shaped [B, L, n_output].

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

def fused_rollout(spec: RENSpec, u: Tensor, x0: Tensor, e: ExplicitREN) -> tuple[Tensor, Tensor]:
    """Run the rollout through the fused-kernel custom ops (autograd-capable).

    Picks the training op (which stores the states and equilibrium activations for the
    analytic BPTT backward) when gradients are live, else the inference op, which keeps no
    intermediates.

    Returns:
        ``(y, x_L)`` with ``y`` shaped ``[B, L, n_output]``.
    """
    params = e.tensors
    fields = (spec.n_state, spec.n_input, spec.n_output, spec.n_nl, spec.act)
    if torch.is_grad_enabled() and any(t.requires_grad for t in (u, x0, *params)):
        y, x_last, _, _ = _ren_rollout_train(u, x0, params, *fields)
        return y, x_last
    return _ren_rollout(u, x0, params, *fields)

lbdn_forward

lbdn_forward(layers: tuple[ExplicitSandwich, ...], h: Tensor, act: Callable, gamma: float = 1.0) -> Tensor

Evaluate an LBDN from its explicit realization alone.

Separate from :class:LBDN so a caller that builds the realization once — a rollout reusing it at every timestep — never pays for the Cayley transforms again, and so the evaluation depends on plain tensors rather than on module state.

Parameters:

Name	Type	Description	Default
layers	tuple[ExplicitSandwich, ...]	per-layer realizations; all but the last are the nonlinear form.	required
h	Tensor	input [..., n_input].	required
act	Callable	the activation, monotone and slope-restricted to [0, 1].	required
gamma	float	Lipschitz bound, applied as √gamma at each end of the stack.	1.0

Source code in tsfast/models/architectures/ren/lbdn.py

def lbdn_forward(layers: tuple[ExplicitSandwich, ...], h: Tensor, act: Callable, gamma: float = 1.0) -> Tensor:
    """Evaluate an LBDN from its explicit realization alone.

    Separate from :class:`LBDN` so a caller that builds the realization once — a rollout
    reusing it at every timestep — never pays for the Cayley transforms again, and so the
    evaluation depends on plain tensors rather than on module state.

    Args:
        layers: per-layer realizations; all but the last are the nonlinear form.
        h: input ``[..., n_input]``.
        act: the activation, monotone and slope-restricted to ``[0, 1]``.
        gamma: Lipschitz bound, applied as ``√gamma`` at each end of the stack.
    """
    scale = math.sqrt(gamma)
    h = scale * h
    for e in layers[:-1]:
        h = _sandwich_forward(e, h, act)
    return _sandwich_forward(layers[-1], scale * h, act)

folded_weights

folded_weights(layers: tuple[ExplicitSandwich, ...]) -> list[Tensor]

The network collapsed to one matrix and one bias per layer, in fused-kernel order.

A hidden layer's √2 A Ψ σ(√2 Ψ⁻¹ B h + c) is V σ(W h + c) with W = √2 Ψ⁻¹ B and V = √2 A Ψ. Every factor there is constant across a rollout, so this folds them once per call — and then folds each V into the following layer's W, since W_{l+1}(V_l a_l) = (W_{l+1} V_l) a_l. What reaches a kernel is one matrix per layer: half the register footprint and, more to the point, half the dependent cross-lane reductions per timestep, which is what a sequential rollout is actually bound by (MATH_R2DN.md §1). Autograd carries ∂L/∂B, ∂L/∂ψ and ∂L/∂A back through the composition.

Returns:

Type	Description
list[Tensor]	[W_0, c_0, ..., W_out, c_out] — two tensors per layer, output layer last.

Source code in tsfast/models/architectures/ren/lbdn.py

def folded_weights(layers: tuple[ExplicitSandwich, ...]) -> list[Tensor]:
    """The network collapsed to one matrix and one bias per layer, in fused-kernel order.

    A hidden layer's ``√2 A Ψ σ(√2 Ψ⁻¹ B h + c)`` is ``V σ(W h + c)`` with ``W = √2 Ψ⁻¹ B``
    and ``V = √2 A Ψ``. Every factor there is constant across a rollout, so this folds them
    once per call — and then folds each ``V`` into the *following* layer's ``W``, since
    ``W_{l+1}(V_l a_l) = (W_{l+1} V_l) a_l``. What reaches a kernel is one matrix per layer:
    half the register footprint and, more to the point, half the dependent cross-lane
    reductions per timestep, which is what a sequential rollout is actually bound by
    (``MATH_R2DN.md`` §1). Autograd carries ``∂L/∂B``, ``∂L/∂ψ`` and ``∂L/∂A`` back through
    the composition.

    Returns:
        ``[W_0, c_0, ..., W_out, c_out]`` — two tensors per layer, output layer last.
    """
    out, pending = [], None
    for e in layers:
        w = e.B if e.A is None else _SQRT2 * e.B / e.psi[:, None]
        out += [w if pending is None else w @ pending, e.bias]
        pending = None if e.A is None else _SQRT2 * e.A * e.psi
    return out