Skip to content

Vectorize structured matrices densely - #253

Open
devmotion wants to merge 3 commits into
mainfrom
dmw/dense-matrix-charts
Open

devmotion wants to merge 3 commits into
mainfrom
dmw/dense-matrix-charts

Conversation

@devmotion

@devmotion devmotion commented Oct 26, 2026 •

Copy link
Copy Markdown
Member

Follow-up to #252.

Motivation

j′vp is only correct if from_vec maps cotangents correctly, and it uses the same from_vec for primals and cotangents. With dense coordinates this holds if from_vec is the orthogonal projection onto the matrices of the given structure. Moreover, j′vp pairs to_vec(f(x)) with to_vec(ȳ), so all matrix types have to use the same coordinates (dense, column-major) for tangents of different types to be compatible (cf. #186). Diagonal, triangular, Symmetric/Hermitian, Adjoint/Transpose and sparse matrices are already vectorized this way, but the following types were not. On 0.12.34:

Type Before
Tridiagonal vectorized by its fields: j′vp with a dense output cotangent throws a DimensionMismatch
SymTridiagonal generic struct fallback: off-diagonals of cotangents are doubled (same as #164), e.g. grad(sum, x) has 2.0 on the off-diagonals; also DimensionMismatch
Bidiagonal not supported: Expected a struct type (the uplo::Char field)
PermutedDimsArray vectorized in the order of the parent: j′vp(fdm, identity, A, PermutedDimsArray(A, (2, 1))) silently returns A' instead of A

Changes

  • Tridiagonal, SymTridiagonal, Bidiagonal and PermutedDimsArray are vectorized densely (in the order of the array itself), and their from_vec projects onto their structure (for SymTridiagonal: the off-diagonals are averaged).
  • Symmetric/Hermitian reuse the from_vec of their parent to project the symmetrized matrix onto the parent's structure. This adds support for e.g. Diagonal, triangular and Tridiagonal parents, which failed before (Fix double counting of Symmetric/Hermitian off-diagonals #252 did not change this). Since this relies on the parent being vectorized densely, an ArgumentError is thrown otherwise (e.g. for custom matrix types vectorized by the generic struct fallback). Parents whose from_vec returns a different type (e.g. views, which are vectorized via copy) yield a Symmetric/Hermitian of that type.
  • Version 0.12.36.

Behaviour changes

The length of to_vec(x) changes for Tridiagonal and SymTridiagonal, and the order changes for PermutedDimsArray, so code that compares values in the vectorized space would be affected. The tests of ChainRulesTestUtils pass locally with this branch.

Tests

  • j′vp/jvp through identity with dense tangents for all four types (real and complex), comparing with the orthogonal projection.
  • The Symmetric double counts #164 testset now covers Diagonal, UpperTriangular, LowerTriangular and Tridiagonal parents of Symmetric/Hermitian.
  • Round trips for the new types/parents, and the ArgumentError.

Longer term, the projections could be replaced by ChainRulesCore.ProjectTo.

🤖 Generated with Claude Code

devmotion and others added 2 commits October 26, 2026 16:54
to_vec(::HermOrSym) rebuilt the matrix from one triangle of the dense vector,
so cotangents computed by j′vp and grad had their off-diagonals doubled.
Rebuilding from the symmetric/hermitian part instead gives the natural
cotangent and still round-trips primals exactly.

Fixes #164.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Tridiagonal, SymTridiagonal, Bidiagonal and PermutedDimsArray are now
vectorized like dense matrices, and their from_vec projects onto their
structure. This makes them compatible with tangents of other types (e.g.
dense output cotangents in j′vp), fixes the double counting of the
off-diagonals of SymTridiagonal (same as #164), fixes the silently
transposed cotangents of PermutedDimsArray, and adds support for
Bidiagonal.

Symmetric/Hermitian reuse the from_vec of their parent to project onto
its structure, which adds support for e.g. Diagonal, triangular and
Tridiagonal parents.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Base automatically changed from dmw/symmetric-to-vec to main October 28, 2026 07:25

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant