Skip to content

Preserve higher-order triangular matmul derivatives via analytic VJP - #152

Open
M-Colley wants to merge 2 commits into
cornellius-gp:mainfrom
M-Colley:codex/triangular-matmul-gradgrad
Open

M-Colley wants to merge 2 commits into
cornellius-gp:mainfrom
M-Colley:codex/triangular-matmul-gradgrad

Conversation

@M-Colley

Copy link
Copy Markdown

Enable higher-order matrix derivatives through triangular matmul by delegating the bilinear VJP to the wrapped operator. The triangular wrapper preserves the representation, so it can retain an existing analytic derivative instead of using the generic detached autograd fallback.

For TriangularLinearOperator(L).matmul(z).square().sum() with constant z, the first derivative currently has requires_grad=False, although dense matmul has a nonzero matrix Hessian. This also loses a term when differentiating Cholesky-root samples twice through their covariance parameters. This is a narrow higher-order-autograd enhancement; it does not assert a library-wide guarantee for arbitrary operators.

Validation: all 130 tests and 12 subtests in the triangular-operator module pass. Coverage includes vector/matrix and upper/lower gradchecks, a direct constant-RHS matrix-Hessian regression on CPU/CUDA, and independent dense comparisons for Cholesky-sample first and second derivatives. Against unchanged main, the two Cholesky cases and four constant-RHS subtests fail. Formatting, lint and diff checks pass.

Prepared with Codex; tests were exercised locally against unchanged upstream main.

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