pyzag.operators

The pyzag.operators subpackage defines the abstract block-operator interfaces the solver runs against, plus the concrete dense backend. New in pyzag 2.0.

Mathematical model

The solver works on chunks of a recursive nonlinear system. Within a chunk of $m$ steps the unknowns are stacked into a block vector

$x = (x_1, \dots, x_m), \qquad x_k \in \mathbb{R}^{n},$

with logical shape (nblk, batch_size, block_size) – axis 0 is the block / time index, axis 1 is the batch, and the last axis holds one block’s entries. A backend is free to store this however it likes (dense tensor, per-group list, sparse) as long as it honours that logical shape.

A block operator $A$ is a linear map defined purely by its action – it need not materialize a matrix. The two required actions are the matrix-vector product and its transpose,

$y = A\,x, \qquad y = A^{\mathsf{T}} x,$

exposed as matvec() and t_matvec(). A solvable block operator additionally solves the block system

$A\,x = b \;\Longrightarrow\; x = A^{-1} b,$

via solve().

For lookback 1 the chunk residual Jacobian is block lower-bidiagonal: with diagonal blocks $A_k = \partial R_k/\partial x_k$ and subdiagonal blocks $B_k = \partial R_k/\partial x_{k-1}$ ,

$(J\,x)_k = A_k\,x_k + B_{k-1}\,x_{k-1},$

and all other blocks are zero. Newton solves $J\,\Delta = R$ each iteration. This bidiagonal solve is provided two ways: Thomas (sequential forward substitution) and parallel cyclic reduction (PCR), which repeatedly halves a power-of-two window by eliminating the odd blocks using $A^{-1}$ -products. The PCR working state is carried opaquely across levels by a PCRState. The BlockJacobian abstraction owns the per-chunk Jacobian and knows how to build both the forward system (for Newton) and the transposed adjoint system (for the reverse/adjoint sweep).

Note

The concrete bidiagonal solve operators that consume these interfaces – BidiagonalForwardOperator and the Thomas / PCR / hybrid factorizations that implement the $J^{-1}$ action described above – live in pyzag.chunktime, since they are solver machinery rather than storage backends.

Abstract interfaces

Every backend implements these ABCs from pyzag.operators.base. The contract is behavioral: a backend may use any storage as long as the logical block ordering and shapes are preserved.

Block vectors

class pyzag.operators.base.BlockVector

Bases: ABC

Abstract interface for a logical packed block vector.

Logical conventions

A block vector consists of nblk logical blocks. Solver code treats block vectors as block-major: axis 0 indexes logical blocks.

For the current bidiagonal solvers, the logical shape is

(nblk, batch_size, block_size)

where:

axis 0 is the logical block / time index
axis 1 is the batch index
axis 2 is the local vector entries for one block

Backend freedom

A backend may store data in any representation it wants (dense, structured, sparse, factored, etc.) as long as:

nblk, batch_size, block_size describe the logical shape

all methods preserve the same logical block ordering

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk: int: Number of logical blocks in this vector.

abstract property batch_size: int: Logical batch size.

abstract property block_size: int: Logical size of one block (last axis).

abstractmethod clone() → BlockVector: Return a safe copy of the vector.

abstractmethod norm(dim: int = -1) → Tensor: Compute the norm along dim. Returns a raw tensor (used for scalar convergence checks; the result has a different shape than a block vector and should not be wrapped).

abstractmethod flatten() → Tensor: Flatten to a raw tensor of shape (batch_size, ndof): batch-major, with every block and state entry for a given batch element concatenated along the last axis.

abstractmethod where(mask: Tensor, other: BlockVector) → BlockVector

Batch-axis conditional combination, torch.where convention: returns self where mask is True, else other.

mask is a raw tensor of shape (batch_size,); backends broadcast it over the block and state axes. Used by the Newton step for partial updates.

abstractmethod scale_batches(factor: Tensor) → BlockVector: Per-batch scalar broadcast: multiply each batch element by the corresponding entry of factor (raw tensor of shape (batch_size,), broadcast over block and state axes). Used by line-search backtracking.

abstractmethod flip(dim: int) → BlockVector: Return a new vector with axis dim reversed.

neg() → BlockVector: Return the negated vector (alias for -self).

abstractmethod classmethod cat(vectors: Sequence[BlockVector], dim: int = 0) → BlockVector: Concatenate a sequence of compatible block vectors along dim.

abstractmethod classmethod zeros_like(other: BlockVector) → BlockVector: Construct a zero-filled block vector with the same shape as other.

Block operators

class pyzag.operators.base.BlockOperator

Bases: ABC

Abstract interface for a logical packed block operator.

Logical conventions

The operator consists of nblk logical blocks. The solver treats block vectors passed to this operator as block-major.

Backend freedom

A backend may store its blocks in any representation it wants as long as:

nblk reports the correct logical number of blocks
batch_size describes the logical block action
all methods preserve the same logical block ordering

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk: int: Number of logical blocks in this operator.

abstract property batch_size: int: Logical batch size expected in block-major vector inputs.

abstractmethod matvec(x: BlockVector) → BlockVector

Apply the operator to a block vector x.

Required:: x.nblk == self.nblk

abstractmethod t_matvec(x: BlockVector) → BlockVector

Apply the transpose of the operator to a block vector x.

Required:: x.nblk == self.nblk

abstractmethod clone() → BlockOperator: Return a safe copy of the operator.

abstractmethod pad_front(n: int = 1) → BlockOperator: Return an operator with n leading dummy logical blocks (creates new data, not a view).

class pyzag.operators.base.SolvableBlockOperator

Bases: BlockOperator

Block operator supporting direct block solves and PCR-based reduction.

abstractmethod solve(rhs: BlockVector) → BlockVector

Solve the block system A x = rhs.

Required:: rhs.nblk == self.nblk

abstractmethod pcr_init(B: BlockOperator, v: BlockVector) → PCRState

Initialise the backend-native PCR working state for a power-of-two window.

self is the diagonal operator (A), B is the subdiagonal (already padded with one dummy leading block so B.nblk == self.nblk), and v is the RHS slice for this window.

The backend allocates internal working tensors (e.g. adds the extra leading dimension used by the Dense cyclic-shift trick) and returns an opaque PCRState.

abstractmethod pcr_reduce_level(state: PCRState, level: int) → PCRState

Apply one PCR level to the working state.

Updates the RHS vector and subdiagonal via A-inverse products, then applies the backend-native cyclic interleaving (as_strided for Dense). Returns the updated PCRState ready for the next level or for pcr_finalize().

level (0-based) lets the backend compute the correct stride pattern without the caller knowing about the internal working shape.

abstractmethod pcr_finalize(state: PCRState) → tuple[BlockOperator, BlockVector]

Extract the reduced (B_red, v_red) from the final PCR state.

Returns a pair with nblk == window_size - 1, suitable for writing back into the full B and v_work arrays in pyzag.chunktime.BidiagonalPCRFactorization.

class pyzag.operators.base.PCRState

Bases: ABC

Opaque PCR working state managed by the backend across levels.

Created by SolvableBlockOperator.pcr_init(), updated by SolvableBlockOperator.pcr_reduce_level(), and consumed by SolvableBlockOperator.pcr_finalize(). Callers in chunktime.py treat this as an opaque handle and never inspect its contents.

Block Jacobians

class pyzag.operators.base.BlockJacobian

Bases: ABC

Abstract per-chunk Jacobian for a recursive nonlinear system.

Logical model

For lookback = 1 (the only case the solver currently supports), the chunk Jacobian represents the linearization of R[k] = f(x[k-1], x[k]) over k = 1..nblk_steps. The two structural pieces are:

the diagonal: dR[k]/dx[k] for k = 1..nblk_steps

the subdiagonal: dR[k]/dx[k-1] for k = 1..nblk_steps

The boundary subdiagonal blocks couple a chunk to its neighbours. In forward time order, sub[0] (k = 1, dR[1]/dx[0]) couples the chunk’s first residual to the lookback / previous chunk’s last state; the remaining subdiagonal blocks are internal to the chunk’s bidiagonal system. The adjoint pass walks time in reverse (see as_adjoint_walk()) and couples to the previously-processed (adjoint-order) chunk through index 0 of the walk-order subdiagonal – which is the original last forward block, not sub[0].

Time order is forward (low index = early time). The adjoint walk methods internalize all reversal – callers use couple_prev_chunk() and must NOT reach into storage to flip, reorder, or pick a boundary block themselves.

Backend freedom

Backends may store the diagonal/subdiagonal in any layout (dense per block, structured/arrowhead, sparse, factored). The contract is purely behavioral.

abstract property device: device: Execution device for torch-backed implementations.

abstract property dtype: dtype: Data type for torch-backed implementations.

abstract property nblk_steps: int: Number of residual rows in the chunk.

abstract property batch_size: int: Logical batch size.

abstract property block_size: int: Per-step state size (the user-facing n).

abstractmethod forward_system(inverse_operator) → BidiagonalForwardOperator

Build the chunk’s forward bidiagonal system, ready for Newton.

Returns a pyzag.chunktime.BidiagonalForwardOperator whose diagonal A has nblk == nblk_steps and whose subdiagonal B has nblk == nblk_steps - 1.

Parameters:: inverse_operator – factory used to build the inverse (e.g. pyzag.chunktime.BidiagonalThomasFactorization).

abstractmethod adjoint_system(inverse_operator)

Build the chunk’s adjoint bidiagonal solve operator, in adjoint-walk order.

Transposes are baked in. The first-row / first+last-col slicing (J[1, 1:].T and J[0, 1:-1].T after flip(1)) is also baked in. Unlike forward_system(), which returns a BidiagonalForwardOperator (so Newton can call .inverse() between iterations), this returns the inverse / solve operator directly: applying .matvec(rhs) on the returned object yields the adjoint solution. This matches how block_update_adjoint consumes the result with a single linear solve per chunk.

This method should be called on a BlockJacobian returned by as_adjoint_walk().

abstractmethod solve_terminal_adjoint(g_terminal: Tensor) → BlockVector

Compute -A_terminal^{-T} @ g_terminal for the very last forward-time step of the trajectory.

Returns a single-block BlockVector (nblk == 1 for lookback == 1; for higher lookback this would be a lookback-block vector). The returned vector must not alias g_terminal.

abstractmethod couple_prev_chunk(a_first: BlockVector) → BlockVector

Compute the inter-chunk adjoint coupling B_boundary^T @ a_first.

a_first is a single-block BlockVector holding the previous chunk’s adjoint tail (in adjoint-walk order). Returns a single-block BlockVector to be subtracted into the current chunk’s RHS first row (nblk == 1 for lookback == 1; lookback-block for higher lookback).

abstractmethod as_adjoint_walk() → BlockJacobian: Return a BlockJacobian whose forward time order is the reverse of this one. Backends are free to implement lazily (e.g. via a flag) so that no copy happens unless storage requires it. Replaces direct flip calls in solver code.

Dense backend

The dense backend is the default implementation and reproduces the pyzag 1.x behavior. Block vectors are stored as (nblk, batch, state) tensors and block operators as (nblk, sbat, sblk, sblk) tensors. Each abstract interface above has a dense counterpart below.

class pyzag.operators.dense.DenseBlockVector(data: Tensor)

Bases: BlockVector

Dense tensor-backed packed block vector.

Parameters:: data (torch.Tensor) – shape (nblk, sbat, sblk)

property device: device: Execution device of the backing tensor.

property dtype: dtype: Data type of the backing tensor.

property nblk: int: Number of logical blocks (axis 0 of data).

property batch_size: int: Logical batch size (axis 1 of data).

property block_size: int: Logical size of one block (last axis of data).

clone() → DenseBlockVector: Return a safe copy backed by a cloned tensor.

norm(dim: int = -1) → Tensor: Compute the norm along dim. Returns a raw tensor (used for scalar convergence checks; not wrapped as a block vector).

flatten() → Tensor: Flatten to a raw (batch_size, nblk * block_size) tensor: transpose the block axis behind the batch axis and flatten the rest, so each batch element’s entries are concatenated along the last axis.

where(mask: Tensor, other: BlockVector) → DenseBlockVector: Batch-axis conditional combination (torch.where convention): return self where mask is True, else other. mask is a (batch_size,) tensor broadcast over the block and state axes.

scale_batches(factor: Tensor) → DenseBlockVector: Multiply each batch element by its own scalar from factor (shape (batch_size,)), reshaped to (1, batch_size, 1, ...) so it broadcasts over the block and state axes rather than the trailing (state) axis. Used by line-search backtracking.

flip(dim: int) → DenseBlockVector: Return a new vector with axis dim reversed.

classmethod cat(vectors: Sequence[BlockVector], dim: int = 0) → DenseBlockVector: Concatenate a sequence of DenseBlockVector along dim.

classmethod zeros(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockVector: Return a zero-filled DenseBlockVector of the given shape.

classmethod zeros_like(other: BlockVector) → DenseBlockVector: Return a zero-filled DenseBlockVector with the same shape as other.

classmethod empty(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockVector: Return an uninitialized DenseBlockVector of the given shape.

class pyzag.operators.dense.DenseBlockOperator(data: Tensor, lu: Tensor | None = None, pivots: Tensor | None = None)

Bases: SolvableBlockOperator

Dense tensor-backed packed block operator.

Implements SolvableBlockOperator (LU-based solve and PCR primitives pcr_init / pcr_reduce_level / pcr_finalize).

The cached LU factorization (self.lu / self.pivots) is materialized lazily on the first solve or pcr_init call. Use factored() to construct with eager factorization for cases where you know LU will be needed (e.g., the diagonal of a Newton system). Slicing, cloning, and in-place assignment preserve the cached LU when present, so chained Thomas-style A[i:i+1].solve(...) calls don’t re-factor every block.

Parameters:

data (torch.Tensor) – shape (nblk, sbat, sblk, sblk).
lu (torch.Tensor, optional) – pre-computed LU factor with the same shape as data.
pivots (torch.Tensor, optional) – pre-computed pivots, shape (nblk, sbat, sblk). Must be paired with lu.

classmethod factored(data: Tensor) → DenseBlockOperator: Construct with eager LU factorization. Use when the caller knows solve or pcr_init will be invoked.

property device: device: Execution device of the backing tensor.

property dtype: dtype: Data type of the backing tensor.

property nblk: int: Number of logical blocks (axis 0 of data).

property batch_size: int: Logical batch size expected in block-major vector inputs.

matvec(x: BlockVector) → DenseBlockVector: Apply the (block-diagonal) operator to x: per block, the matrix-vector product A_k x_k. Requires x.nblk == self.nblk.

t_matvec(x: BlockVector) → DenseBlockVector: Apply the transpose of the operator to x (per block, A_k^T x_k). Requires x.nblk == self.nblk.

solve(rhs: BlockVector) → DenseBlockVector: Solve the block system A x = rhs per block via the cached LU factorization (materialized lazily on first call). Requires rhs.nblk == self.nblk.

clone() → DenseBlockOperator: Return a safe copy, cloning the cached LU factors if present.

pad_front(n: int = 1) → DenseBlockOperator: Return an operator with n leading zero (dummy) blocks prepended. Creates new data rather than a view; the result is un-factored since padding changes the per-block layout.

pcr_init(B: BlockOperator, v: BlockVector) → DensePCRState

Initialise Dense PCR working state for a power-of-two window.

The state owns its b/v working buffers (cloned here), so the in-place reductions in pcr_reduce_level() never mutate the caller’s operator/vector – honouring the base-class contract that the backend allocates internal working tensors.

pcr_reduce_level(state: PCRState, level: int) → DensePCRState: Apply one Dense PCR level: update v and B, then cyclic-shift all four tensors.

pcr_finalize(state: PCRState) → tuple[DenseBlockOperator, DenseBlockVector]: Extract (B_red, v_red) with nblk = window_size - 1 from the final PCR state.

classmethod identity(nblk: int, batch_size: int, block_size: int, dtype: dtype, device: device) → DenseBlockOperator: Return an identity DenseBlockOperator of the given shape.

classmethod from_diagonal(data: Tensor) → DenseBlockOperator: Return a DenseBlockOperator built from per-block diagonal data.

class pyzag.operators.dense.DensePCRState(lu: Tensor, pivots: Tensor, b: Tensor, v: Tensor)

Bases: PCRState

Dense backend PCR working state.

Holds the four working tensors used by the Dense cyclic-shift PCR kernel. Each tensor has an extra leading dimension prepended by pcr_init() so that _dense_pcr_cyclic_shift() can double it at each level.

lu

shape (1, nblk, sbat, sblk, sblk) initially.

Type:: torch.Tensor

pivots

shape (1, nblk, sbat, sblk) initially.

Type:: torch.Tensor

b

shape (1, nblk, sbat, sblk, sblk) initially.

Type:: torch.Tensor

v

shape (1, nblk, sbat, sblk, 1) initially.

Type:: torch.Tensor

class pyzag.operators.dense.DenseBlockJacobian(diag: Tensor, sub: Tensor, _reversed: bool = False)

Bases: BlockJacobian

Dense tensor-backed per-chunk Jacobian.

Storage is canonical (forward-time order). The as_adjoint_walk() method returns a sibling instance that shares the same underlying tensors and only flips a private flag; the four adjoint methods read the data with appropriate slicing/indexing under that flag, so a physical flip never materializes.

Parameters:

diag (torch.Tensor) – per-step diagonal blocks dR[k]/dx[k], shape (nblk_steps, batch, n, n).
sub (torch.Tensor) – per-step subdiagonal blocks dR[k]/dx[k-1], shape (nblk_steps, batch, n, n). The first block sub[0] is the inter-chunk boundary coupling to the lookback / previous chunk’s last state.

classmethod from_stacked(J: Tensor) → DenseBlockJacobian: Construct from a stacked tensor of shape (2, nblk_steps, batch, n, n) where J[1] is the diagonal and J[0] is the subdiagonal (with J[0, 0] being the boundary block).

property device: device: Execution device of the backing tensors.

property dtype: dtype: Data type of the backing tensors.

property nblk_steps: int: Number of residual rows in the chunk (axis 0 of diag).

property batch_size: int: Logical batch size (axis 1 of diag).

property block_size: int: Per-step state size (last axis of diag).

forward_system(inverse_operator): Build the chunk’s forward bidiagonal system for Newton: diagonal A = the (factored) per-step diag blocks and subdiagonal B = sub[1:], wrapped in a BidiagonalForwardOperator. Must be called on a forward-walk Jacobian.

adjoint_system(inverse_operator): Build the chunk’s adjoint bidiagonal solve operator, in adjoint-walk order with transposes baked in: drop the first walked diagonal block (handled by the terminal seed) and the first/last walked subdiagonal blocks. Must be called on the walk-order Jacobian from as_adjoint_walk(); applying .matvec(rhs) yields the adjoint solution.

solve_terminal_adjoint(g_terminal: Tensor) → DenseBlockVector: Compute -A_terminal^{-T} @ g_terminal for the very last forward-time step, returned as a single-block DenseBlockVector that does not alias g_terminal.

couple_prev_chunk(a_first: BlockVector) → DenseBlockVector: Compute the inter-chunk adjoint coupling B_boundary^T @ a_first, where a_first is the previous chunk’s adjoint tail (single-block, in walk order). Returns a single-block DenseBlockVector to subtract into the current chunk’s first-row RHS.

as_adjoint_walk() → DenseBlockJacobian: Return a sibling Jacobian whose forward time order is reversed. This shares the same underlying tensors and only toggles a private flag, so no physical flip materializes.

Dense helpers

pyzag.operators.dense.batch_lu_solve(lu: Tensor, pivots: Tensor, rhs: Tensor) → Tensor: Batched version of torch.linalg.lu_solve that accepts separate LU and pivot tensors.