pyzag abstraction and neml2 v3 compatibility

[hdt5kt]

this is compatible with a recent PR in NEML2 v3

Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com

Release 2.0.0: block-operator abstraction + NEML2 backend

- operators/: BlockVector/BlockOperator/SolvableBlockOperator/BlockJacobian/ PCRState ABCs + dense backend (default, 1.x-equivalent) + a pure-torch NEML2 backend (no `import neml2`).
- chunktime/nonlinear/ode rewritten against the abstraction.
- Correctness fixes: y0 adjoint gradient now includes the objective's direct term; IC-only adjoint (empty params) no longer crashes; single-chunk reverse() guarded; PCR working buffers no longer alias caller tensors; NEML2 vector norm -> combined L2 (matches dense/NEML2); LU cache preserved on slice; factory lookback/wrapper are abstract.
- Tests: dense+NEML2 block operator/vector/jacobian, non-power-of-two PCR, PCR/Hybrid adjoint sweep, IC adjoint; dropped example-driven smoke tests.
- Docs: operators API page + MIGRATION.md.
- Examples: curated standalone examples (dense: mass_damper_spring, neuron, linear_network; NEML2: general, crystal, calibration); dev scripts moved to examples/dev (gitignored).

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-07-27 21:54 UTC

[Copilot]

Pull request overview

This PR is a major 2.0.0 refactor that introduces a backend-agnostic block-operator abstraction layer and adds a pure-torch NEML2-compatible backend, while updating ODE integration and solver pathways to operate through the new BlockVector / BlockJacobian interfaces.

Changes:

• Add pyzag.operators abstraction layer (vector/operator/jacobian + PCR state) and implement dense + NEML2 backends against it.
• Rewrite ODE integration as NonlinearFunctionOperatorFactory implementations returning backend-typed BlockJacobian objects.
• Refresh/expand the test suite, docs, examples, and packaging metadata for the 2.0 API and new backends.

Reviewed changes

Copilot reviewed 42 out of 46 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
test/test_ode.py	Updates ODE tests to use evaluate_raw and assert DenseBlockJacobian shape/metadata.
test/test_initial_conditions.py	Fixes adjoint wrapper test driver to return adjoint result so .backward() flows through wrapper.
test/test_chunktime_dense.py	Refactors chunktime dense tests to use dense block operator/vector backends.
test/test_block_vector_neml2.py	Adds comprehensive tests for NEML2BlockVector (single + multi-group layouts).
test/test_block_vector_dense.py	Adds tests for DenseBlockVector contract/operations.
test/test_block_operator_neml2.py	Adds tests for NEML2 solvable block operator algebra + solve/transpose semantics.
test/test_block_operator_dense.py	Adds tests for dense block operator algebra, PCR primitives, and dense block jacobian behavior.
test/test_block_jacobian_to_dense.py	Adds equivalence tests between NEML2 jacobian materialization and dense jacobian solves.
test/test_block_jacobian_contract.py	Adds contract tests ensuring solver core uses only abstract interfaces (no .data[ reach-through).
test/test_adjoint.py	Extends adjoint gradient checks to sweep Thomas/PCR/Hybrid and ports nonlinear factory to new interface.
pyzag/stochastic.py	Adds typing/cleanup and removes .data usage in parameter mapping.
pyzag/reparametrization.py	Adds typing and replaces unsafe .data mutation with copy_ under no_grad.
pyzag/operators/neml2/_wrapper.py	Introduces NEML2 backend wrapper bridging flat tensors and structured vectors/jacobians.
pyzag/operators/neml2/_vector.py	Implements NEML2BlockVector operations and layout-aware factories.
pyzag/operators/neml2/_jacobian.py	Implements NEML2BlockJacobian (forward/adjoint systems, coupling, dense materialization).
pyzag/operators/neml2/_containers.py	Adds pure-torch NEML2-mirror container types (layout/tensor/matrix/vector).
pyzag/operators/neml2/init.py	Exposes NEML2 backend public surface + design notes.
pyzag/operators/base.py	Adds the abstract BlockVector/BlockOperator/BlockJacobian/PCRState interfaces.
pyzag/operators/init.py	Introduces operators subpackage marker/init.
pyzag/ode.py	Rewrites ODE integration as factory/operator pattern returning BlockJacobian via ODEWrapper.
pyproject.toml	Bumps version to 2.0.0 and adjusts setuptools package discovery include/exclude.
MIGRATION.md	Adds 1.x → 2.0 migration guide documenting API breaks and replacements.
examples/neuron/neuron.py	Adds standalone dense example for neuron network ODE + adjoint/autograd/FD comparison.
examples/neml2_general/neml2_general.py	Adds NEML2-driven example using NEML2-side adapter to run chunked solve.
examples/neml2_general/model.i	Adds NEML2 model input file used by the general example.
examples/neml2_crystal/neml2_crystal.py	Adds multi-grain Taylor crystal example (BLOCK+DENSE Schur backend usage).
examples/neml2_crystal/model_mixed.i	Adds NEML2 mixed-control Taylor model used by crystal example.
examples/neml2_calibration/neml2_calibration.py	Adds end-to-end calibration example using adjoint gradients + reparametrization.
examples/neml2_calibration/model_mixed_calibration.i	Adds calibration-specific NEML2 model input with tunable parameters.
examples/mass_damper_spring/mass_damper_spring.py	Adds standalone dense mass-damper-spring example + gradient comparisons.
examples/linear_network/linear_network.py	Adds standalone dense neural-network ODE example + gradient comparisons.
docs/api/operators.rst	Adds API docs page for operators abstraction and dense backend.
docs/api/api.rst	Wires operators page into API doc index.
.pylintrc	Updates pylint limits/disabled checks for new code scale/structure.
.gitignore	Updates ignore rules for new example outputs/dev artifacts and curated examples.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[reverendbedford]

Also:

1. Just merge the requirements.txt and dev-requirements.txt into the pyproject.toml file and get rid of them. Gary has a nice example in NEML2 of how to manage both "regular" and "dev" requirements.
2. While annoying, you need to repeat the docstrings in the actual implementations of the operators.

[Copilot]

⚠️ Not ready to approve

PCR/Hybrid solvers aren’t currently guarded against autograd use (despite intended no_grad semantics) and LU factorization ignores lu_factor_ex error flags, both of which can cause silent correctness failures.

Review details

Comments suppressed due to low confidence (2)

pyzag/chunktime.py:420

• BidiagonalHybridFactorizationImpl.matvec() has the same differentiability issue as PCR: it performs PCR reduction with in-place/as_strided updates, but does not enforce the intended torch.no_grad execution. Add a grad-context guard so callers don’t accidentally backprop through the hybrid solver.

    def matvec(self, v: BlockVector) -> BlockVector:
        """Apply the hybrid PCR/Thomas factorization."""
        B = self.B.pad_front(1)
        v_work = v.clone()

pyzag/operators/dense.py:96

• Same issue inside the per-matrix CUDA fallback loop: lu_factor_ex's info is ignored for each block, so singular blocks won’t be detected early. Enable error checking so a bad factorization raises at the source.

    for i in range(flat.shape[0]):
        lu_i, piv_i, _ = torch.linalg.lu_factor_ex(flat[i])
        lus_flat[i] = lu_i
        pivs_flat[i] = piv_i

• Files reviewed: 33/38 changed files
• Comments generated: 2
• Review effort level: Low

Note

Your feedback helps us improve the quality of this feature.
Please use 👍 or 👎 to tell us whether this assessment is correct.

[reverendbedford]

I looked through and was happy but copilot found two things. I don't think the first is a real issue, but the second might be something to consider. Either way, take a look and choose something to do and then we'll merge.