Least Squares AD Notes (`lstsq`)

Public contract

lstsq(a, b) returns

solution = pinv(a) @ b
residuals = ||a @ solution - b||_F^2 per right-hand side when m > n and the solve is full-rank
residuals = [] otherwise

The auxiliary metadata rank and singular_values are not differentiated.

The first-order rules are expressed in terms of the pseudoinverse:

For x = pinv(A) b,

dx = d(pinv(A))\, b + pinv(A)\, db

In the implementation, d(pinv(A)) is provided by pinv_frule.

Let

r = A x - b, \qquad dr = dA\,x - db

Then, for each right-hand side,

d\,\mathrm{residuals} = 2 \sum \mathrm{Re}(r \odot \overline{dr})

For the current real-valued lstsq AD path, this reduces to

d\,\mathrm{residuals} = 2 \sum r \odot dr

Let gx be the cotangent of solution. Since x = pinv(A) b,

This is the path used by the implementation.

Let gr be the cotangent of the summary residual outputs, broadcast per RHS. Then

\bar{A}_{res} = 2 (gr \odot r)\, x^H

\bar{b}_{res} = -2 (gr \odot r)

The full VJP is the sum of the solution-path and residual-summary-path contributions.