[devmotion]

Follow-up to #252 (stacked on it; I'll retarget to main once #252 is merged).

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