Abstraction for the pyzag input tensor#21
Conversation
hdt5kt
changed the title
|
@reverendbedford here is the PR, the failed is in the documentation for generating github page. Feel free to change anything as you see fit. |
There was a problem hiding this comment.
Pull request overview
This PR introduces a packed “block operator” abstraction to decouple chunktime/nonlinear from assuming raw dense tensors for bidiagonal block systems, while providing a dense-backed implementation to preserve prior behavior.
Changes:
- Added
BlockOperator/SolvableBlockOperator/BlockOperatorBuilderabstractions and a dense implementation (DenseBlockOperator,DenseBlockLUFactorizedOperator,DenseBlockOperatorBuilder). - Updated
chunktimebidiagonal operators/factorizations andnonlinearto construct and operate on these block-operator objects. - Updated existing tests to use the new abstractions and added a new dense block-operator algebra test suite.
Reviewed changes
Copilot reviewed 7 out of 10 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
pyzag/operators/base.py |
Introduces abstract interfaces for packed block operators and builders. |
pyzag/operators/dense.py |
Adds dense tensor-backed operator implementations + LU-factorized variant + builder. |
pyzag/chunktime.py |
Migrates bidiagonal operator/factorizations to consume BlockOperator abstractions; updates PCR/Thomas paths. |
pyzag/nonlinear.py |
Uses a BlockOperatorBuilder to wrap Jacobian blocks before constructing bidiagonal operators/solvers. |
test/test_chunktime.py |
Updates chunktime tests to use dense operator classes and a new dense-matrix reference path. |
test/test_block_operator_dense.py |
New tests validating dense block-operator algebra and builder behavior. |
test/test_initial_conditions.py |
Removes a comment line near the manual seed call. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from .base import BlockOperator, SolvableBlockOperator, BlockOperatorBuilder | ||
| import torch | ||
|
|
There was a problem hiding this comment.
pyzag/operators is introduced as a new subpackage, but the directory currently has no __init__.py and pyproject.toml does not configure setuptools to discover namespace packages. This can lead to wheels/sdists missing pyzag.operators.* and runtime import failures after installation. Consider adding pyzag/operators/__init__.py (simplest) or explicitly configuring namespace package discovery in the packaging config.
| self.data = data | ||
| self.lu, self.pivots, _ = torch.linalg.lu_factor_ex(data) | ||
|
|
||
| # lazy transpose factorization | ||
| self.t_lu = None | ||
| self.t_pivots = None | ||
|
|
||
| def _ensure_transpose_lu(self): | ||
| if self.t_lu is None or self.t_pivots is None: | ||
| self.t_lu, self.t_pivots, _ = torch.linalg.lu_factor_ex( | ||
| self.data.transpose(-1, -2) | ||
| ) | ||
|
|
There was a problem hiding this comment.
torch.linalg.lu_factor_ex returns an info tensor indicating singular/ill-conditioned blocks. The current code discards it (_) and proceeds, which can silently produce incorrect results or NaNs during lu_solve. Please capture and validate info (e.g., raise if any nonzero) for both the forward and transpose factorizations.
|
|
||
| def empty_like(self, nblk): | ||
| shape = (nblk,) + self.data.shape[1:] | ||
| return DenseBlockLUFactorizedOperator( |
There was a problem hiding this comment.
DenseBlockLUFactorizedOperator.empty_like() currently builds a new factorized operator from torch.empty(...), which immediately LU-factorizes uninitialized memory. Besides being expensive, this can yield invalid LU/pivots and hard-to-debug failures if the caller expects to fill data later. Consider either (a) returning a non-factorized DenseBlockOperator here, (b) allocating a safe default (e.g., identity blocks) before factorization, or (c) raising NotImplementedError to avoid unsafe use.
| return DenseBlockLUFactorizedOperator( | |
| # Return a non-factorized operator to avoid LU-factorizing uninitialized memory. | |
| return DenseBlockOperator( |
| pass | ||
|
|
||
| @abstractmethod | ||
| def empty_like(self, nlbk): |
There was a problem hiding this comment.
The abstract method empty_like uses the parameter name nlbk, but implementations use nblk. While this won’t break positional calls, it’s inconsistent and can confuse readers/type checkers. Rename the parameter to nblk here to match the rest of the API.
| def empty_like(self, nlbk): | |
| def empty_like(self, nblk): |
|
|
||
| if B.nblk != A.nblk - 1: | ||
| raise ValueError("B must have nblk = A.nblk - 1.") | ||
|
|
There was a problem hiding this comment.
BidiagonalOperator.__init__ validates block counts, but it doesn’t validate that A and B agree on batch_size, dtype/device, and block_shape. Mismatches will produce hard-to-diagnose shape/device errors later in matvec/solve. Consider adding explicit checks here (e.g., A.batch_size == B.batch_size and A.block_shape == B.block_shape).
| # Ensure that A and B are fully compatible; otherwise subsequent | |
| # matvec/solve operations may fail with obscure shape/device errors. | |
| if getattr(A, "batch_size", None) != getattr(B, "batch_size", None): | |
| raise ValueError( | |
| f"A and B must have the same batch_size " | |
| f"(got A.batch_size={getattr(A, 'batch_size', None)}, " | |
| f"B.batch_size={getattr(B, 'batch_size', None)})." | |
| ) | |
| if getattr(A, "dtype", None) != getattr(B, "dtype", None): | |
| raise ValueError( | |
| f"A and B must have the same dtype " | |
| f"(got A.dtype={getattr(A, 'dtype', None)}, " | |
| f"B.dtype={getattr(B, 'dtype', None)})." | |
| ) | |
| if getattr(A, "device", None) != getattr(B, "device", None): | |
| raise ValueError( | |
| f"A and B must be on the same device " | |
| f"(got A.device={getattr(A, 'device', None)}, " | |
| f"B.device={getattr(B, 'device', None)})." | |
| ) | |
| if getattr(A, "block_shape", None) != getattr(B, "block_shape", None): | |
| raise ValueError( | |
| "A and B must have the same block_shape " | |
| f"(got A.block_shape={getattr(A, 'block_shape', None)}, " | |
| f"B.block_shape={getattr(B, 'block_shape', None)})." | |
| ) |
| def __init__(self, A, B, *args, **kwargs): | ||
| if isinstance(A, DenseBlockOperator): | ||
| A = DenseBlockLUFactorizedOperator(A.data) | ||
| super().__init__(A, B, *args, **kwargs) | ||
|
|
||
| Args: | ||
| v (torch.tensor): tensor of shape (nblk, sbat, sblk) | ||
| """ | ||
| return thomas_solve(self.lu, self.pivots, self.B, v) | ||
| def matvec(self, v): | ||
| """call the thomas_solve""" | ||
| return thomas_solve(self.A, self.B, v) |
There was a problem hiding this comment.
BidiagonalThomasFactorization ultimately calls A.solve_lower_bidiagonal(...), which requires that A support solve/t_solve (i.e., be a SolvableBlockOperator). Currently only the DenseBlockOperator case is converted; other BlockOperator implementations would fail at runtime. Consider validating A is solvable (or converting via a builder/factory) and raising a clear error if not.
| def _gen_operators(self): | ||
| self.blk_A = torch.rand(self.nblk, self.sbat, self.sblk, self.sblk) | ||
| self.blk_B = ( | ||
| torch.rand(self.nblk - 1, self.sbat, self.sblk, self.sblk) / 10 | ||
| ) # Diagonal dominance | ||
| self.blk_B = torch.rand(self.nblk - 1, self.sbat, self.sblk, self.sblk) / 10 | ||
|
|
||
| # Diagonal blocks must be solvable operators | ||
| self.Aop = DenseBlockLUFactorizedOperator(self.blk_A) | ||
| self.Bop = DenseBlockOperator(self.blk_B) |
There was a problem hiding this comment.
self.blk_A is generated as a fully random dense matrix, but the test now factorizes it (DenseBlockLUFactorizedOperator(self.blk_A)), which can become flaky if any block is singular/ill-conditioned. To make the test deterministic, consider constructing blk_A to be well-conditioned/invertible (e.g., add a scaled identity per block, similar to test_block_operator_dense.py).
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 10 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| from pyzag import chunktime | ||
|
|
||
| from pyzag.operators.dense import DenseBlockOperatorBuilder | ||
|
|
There was a problem hiding this comment.
pyzag/operators is introduced as a new subpackage, but the repository currently has no pyzag/operators/__init__.py. With the current pyproject.toml (setuptools default package discovery), subdirectories without __init__.py typically aren’t included in built distributions, which would make this import fail in an installed wheel/sdist. Add pyzag/operators/__init__.py (or explicitly configure namespace package discovery in packaging config) to ensure pyzag.operators.* is shipped and importable.
| return A_ops, B_ops | ||
|
|
||
| def make_adjoint_blocks(self, J): | ||
| A_ops = DenseBlockOperator(J[1, 1:].transpose(-1, -2)) |
There was a problem hiding this comment.
make_adjoint_blocks() returns A_ops as a plain DenseBlockOperator, but the adjoint path calls self.direct_solve_operator(A_ops, B_ops). If direct_solve_operator is BidiagonalPCRFactorization/hybrid, it requires A to implement PCRFactorizedDiagonalOps and will raise a TypeError. Consider returning a DenseBlockLUFactorizedOperator for the adjoint diagonal blocks (factorize J[1, 1:].T) so PCR/hybrid methods work for adjoint solves too.
| A_ops = DenseBlockOperator(J[1, 1:].transpose(-1, -2)) | |
| A_ops = DenseBlockLUFactorizedOperator(J[1, 1:].transpose(-1, -2)) |
| shape = (nblk,) + self.data.shape[1:] | ||
| return DenseBlockLUFactorizedOperator( | ||
| torch.empty(shape, dtype=self.dtype, device=self.device) | ||
| ) |
There was a problem hiding this comment.
DenseBlockLUFactorizedOperator.empty_like() currently constructs a new DenseBlockLUFactorizedOperator(torch.empty(...)), which immediately runs lu_factor_ex on uninitialized memory. That’s unnecessarily expensive and can yield meaningless LU/pivots. Prefer allocating data/lu/pivots tensors directly and using from_factored(...) (or a dedicated constructor) so empty_like is a cheap allocation-only operation.
| shape = (nblk,) + self.data.shape[1:] | |
| return DenseBlockLUFactorizedOperator( | |
| torch.empty(shape, dtype=self.dtype, device=self.device) | |
| ) | |
| data_shape = (nblk,) + self.data.shape[1:] | |
| lu_shape = (nblk,) + self.lu.shape[1:] | |
| pivots_shape = (nblk,) + self.pivots.shape[1:] | |
| data = torch.empty(data_shape, dtype=self.dtype, device=self.device) | |
| lu = torch.empty(lu_shape, dtype=self.lu.dtype, device=self.lu.device) | |
| pivots = torch.empty( | |
| pivots_shape, dtype=self.pivots.dtype, device=self.pivots.device | |
| ) | |
| return DenseBlockLUFactorizedOperator.from_factored(data, lu, pivots) |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 11 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| raise ValueError("DenseBlockLUFactorizedOperator requires square blocks.") | ||
|
|
||
| self.data = data | ||
| self.lu, self.pivots, _ = torch.linalg.lu_factor_ex(data) |
There was a problem hiding this comment.
torch.linalg.lu_factor_ex returns an info tensor indicating singular/failed factorizations. The current implementation ignores this, so singular blocks can silently propagate and cause confusing downstream failures (or NaNs) during solves. Consider checking info and raising a clear error if any entry is nonzero (or optionally allowing the caller to decide how to handle singular blocks).
| self.lu, self.pivots, _ = torch.linalg.lu_factor_ex(data) | |
| self.lu, self.pivots, info = torch.linalg.lu_factor_ex(data) | |
| if torch.any(info != 0): | |
| bad_blocks = torch.nonzero(info != 0, as_tuple=False).flatten().tolist() | |
| raise RuntimeError( | |
| "LU factorization failed for one or more dense blocks " | |
| f"(info={info.tolist()}, failing block indices={bad_blocks}). " | |
| "This usually indicates a singular block or another factorization failure." | |
| ) |
| # this is to make sure it is backward compatible | ||
| self.block_operator_builder = ( | ||
| DenseBlockOperatorBuilder() | ||
| if block_operator_builder is None | ||
| else block_operator_builder | ||
| ) |
There was a problem hiding this comment.
The comment says this is “backward compatible”, but the solver now always wraps Jacobian blocks into BlockOperator instances and BidiagonalForwardOperator now enforces BlockOperator inputs. Any existing direct_solve_operator implementation that previously accepted raw tensors will break. Either provide an adapter path (e.g., accept tensor inputs and wrap them with the default builder) or update the API/docs/comment to clearly state the new required signature (direct_solve_operator(A_ops, B_ops)).
| for the data in the corresponding | ||
| tensor. These values d can | ||
| range from -(n-1) to (n-1) | ||
| now handles abstractions |
There was a problem hiding this comment.
The docstring "now handles abstractions" is misleading here: SquareBatchedBlockDiagonalMatrix still assumes data entries are raw tensors (uses .shape, .dtype, .device) and won’t accept BlockOperator instances. Please either update the implementation to accept the new operator abstractions or adjust the docstring to reflect the actual expected inputs (tensors).
| now handles abstractions | |
| Square batched block-diagonal matrix backed by raw tensor blocks. | |
| The entries in ``data`` are expected to be tensors with ``shape``, | |
| ``dtype``, and ``device`` attributes. This class does not currently | |
| accept generic ``BlockOperator`` abstractions as block entries. |
| storing the nblk main diagonal blocks | ||
| B (torch.tensor): tensor of shape (nblk-1,sbat,sblk,sblk) | ||
| storing the nblk-1 off diagonal blocks | ||
| This is now handled abstractions. |
There was a problem hiding this comment.
The class docstring is grammatically unclear ("This is now handled abstractions"). Consider updating it to a clear sentence (e.g., "This is now handled by block-operator abstractions.") so the intent is understandable when browsing generated docs.
| This is now handled abstractions. | |
| This is now handled by block-operator abstractions. |
…inal behaviour, remove the generic path, remove unnecessary operators.
reverendbedford
force-pushed
the
development
branch
from
5d75737 to
e512d9b
Compare
reverendbedford
left a comment
reverendbedford
left a comment
There was a problem hiding this comment.
Apart from the specific comments:
- You are missing a lot of "proper" docstrings, including an "Args", "Keyword Args", "Returns" description of what the method/function does. Not including these will break the documentation.
- You should change things around so that conceptually
NonlinearRecursiveFunctionreturns the operators instead of the block Jacobians (rather than set things up inRecursiveNonlinearEquationSolverwith the "builder" class. However, to avoid breaking backward compatibility you can instead add a new parent class on top ofNonlinearRecursiveFunctioncalledOperatorJacobianNonlinearRecursiveFunctionor something, and make sureRecursiveNonlinearEquationSolverworks with this class. Then redefineNonlinearRecursiveFunctionas a child of that new base class, where the forward operator first calls the actual torch function and then wraps the Jacobian with the dense operators. - Consider some of the copilot requests, particularly the ones fixing your spelling errors and suggesting some additional error checking.
There was a problem hiding this comment.
There must be a way to better organize the class structure so that DenseBlockOperator and DenseBlockLUFactorizedOperator don't have all the copy-pasted code. Maybe a DenseOperator class you can mixin or similar.
|
|
||
| import torch | ||
|
|
||
| # Ensure test consistency |
There was a problem hiding this comment.
Uhh, why remove the comment?
| storing the nblk main diagonal blocks | ||
| B (torch.tensor): tensor of shape (nblk-1,sbat,sblk,sblk) | ||
| storing the nblk-1 off diagonal blocks | ||
| This is now handled abstractions. |
There was a problem hiding this comment.
You need a better comment explaining what this class now does.
| return thomas_solve(self.A, self.B, v) | ||
|
|
||
|
|
||
| ########## NEW PCR ################# |
There was a problem hiding this comment.
We don't want these comments in the final PR.
Instead of (nblk, nbatch, nstate, nstate), the abstraction is now (nblk, xxxxx).