[heydtn]

I've been procrastinating opening this due to the sheer amount of details in it (apologies in advance) and how I haven't had the time. I decided I would open this with a (mostly) reviewed PR message so that folks could take a look at it and see the data (even if it is just a wall of text) and provide any thoughts.

Problem

Nx.LinAlg.invert/1 computed the determinant of the input, returned NaN
when the determinant's magnitude was at most 1000 times the smallest
positive normal number of the input type, and otherwise divided each row
by its largest magnitude and called solve/2 with the identity scaled the
same way. Three things were wrong with that implementation:

1. The determinant threshold measured the size of the entries, not
   singularity. A well-conditioned matrix with small entries, such as
   [[0.2, 0], [0, 0.2]] in f16, has a determinant below the threshold
   and came back NaN.
2. The same threshold missed most singular matrices, because the computed
   determinant of an exactly singular matrix usually rounds to a small
   nonzero value far above the threshold. Of 300 random exactly singular
   8x8 integer matrices in f32, 25 were caught, 11 as NaN and 14 as an
   ArgumentError from solve/2 meeting an exactly zero pivot; of 300 at
   32x32, 2. The rest came back as large finite numbers.
3. With vectorized input, if in defn runs both branches and selects per
   element, so a batch with one singular matrix still ran solve/2 on
   that matrix. On the default backend and on Torchx, solve/2 raised
   ArgumentError for that matrix, and the raise failed the whole batch.

Changes

invert/1 now computes an LU factorization of the row-scaled input,
tests each pivot of U for singularity, and finishes with two triangular
solves. It no longer computes a determinant or calls solve/2.

Row scaling. The input is cast to floating point (integer input
becomes f32, as on main) and each row is multiplied by the power of two
nearest the reciprocal of its largest magnitude, so every nonzero row has
largest magnitude between 1/sqrt(2) and sqrt(2). Zero rows are left
alone. The exponent is clipped to the normal range of the type, because
some backends flush a subnormal scale factor to zero. Multiplying by a
power of two is exact, so an exactly singular matrix stays exactly
singular after scaling; LAPACK's equilibration routine xGEEQUB scales by
powers of the radix for the same reason [1]. Magnitudes are taken after
the cast to floating point, since the absolute value of a signed integer
type's minimum wraps around in that type.

Singularity test. After {p, l, u} = Nx.LinAlg.lu(scaled), pivot i
is |u[i][i]|, and the elimination formed it by subtracting the products
l[i][j] * u[j][i] for every j < i. The matrix counts as singular when
for some i

|u[i][i]| <= multiplier(type, n) * epsilon(type) * sum over j < i of |l[i][j]| * |u[j][i]|

where n is the matrix order, epsilon(type) is the machine epsilon of
the type and multiplier is the new Nx.Type.rounding_error_multiplier/2
described below. The right-hand side is the rounding error the pivot's own
computation could have accumulated: the componentwise backward error bound
for LU factorization, |A - LU| <= n * u * |L| * |U| with u the unit
roundoff [2], applied to one entry. A pivot no larger than that bound is
indistinguishable from a zero that rounding perturbed. The pivot itself is
not among the subtracted products, so a pivot nothing was subtracted from
is flagged only when it is exactly zero; diagonal and upper triangular
matrices, whose elimination subtracts nothing, therefore pass at any size
and tolerance. A pivot that overflowed to infinity is flagged too.

Why a per-pivot test. Two alternatives were measured against it on
random singular and nonsingular matrices in f32, f16 and bf16:

• Ratio of smallest to largest pivot, rejecting when the smallest
  pivot is at most 16 * n * epsilon times the largest. It rejects graded
  matrices, that is, matrices whose rows differ widely in magnitude, such
  as [[1, 0], [1, 1e-6]], which invert accurately. In bf16 it rejects
  most matrices of size 4 and up.
• Reciprocal condition estimate. rcond is the estimate from [8], as
  LAPACK's xGECON computes it, with rcond < epsilon as the rejection
  rule. It catches every singular matrix in the measurement but also
  rejects the whole graded family in f32 at 32x32 (300 of 900 nonsingular
  matrices) and about 40% of nonsingular bf16 matrices at 8x8 and 16x16.
  It answers a different question: whether the inverse is sensitive to
  perturbations of the data, not whether the computed factorization found
  a singular matrix. The invert/1 docs promise NaN for singular
  matrices, not for ill-conditioned ones. NumPy, PyTorch and Julia also
  signal only for singular matrices, never for ill-conditioned ones,
  though they detect singularity by LAPACK's exactly zero pivot, a
  stricter test than this one; MATLAB estimates the condition number and
  warns. The estimate is also more expensive: an iteration of extra
  triangular solves per matrix, where the per-pivot test is one
  elementwise product and a row sum.

The per-pivot test flags only pivots that rounding could have produced
from zero, which is what the docs promise. A condition number estimate
would be a useful separate function.

Batches. Only the singular matrices in a batch come back NaN. Because
some backends' triangular solves raise on a singular factor, the L and U
of a flagged matrix are swapped for the identity before the solves, and
the NaN replaces that matrix's result.

Nx.Type.rounding_error_multiplier/2 (new, @doc false) returns
2 * n for 32-bit and wider types, min(2 * n, 4 * sqrt(n)) for f16 and
min(2 * n, 4 * sqrt(n), 12) for bf16. The multiplier has to sit above
the rounding noise an exactly singular matrix leaves in its zero pivot,
or singular matrices slip through, and below the pivots of usable
nonsingular matrices, or those are rejected. 2 * n * epsilon is four
times the n * u bound, since epsilon is 2 * u; the extra factor
covers rounding errors in earlier pivots propagating into later ones, and
in f32 and f64 it costs nothing, since no nonsingular matrix in any test
below is rejected. In the 16-bit types the worst-case bound is unusable:
2 * n * epsilon reaches 1 at n = 64 in bf16, where every pivot would be
flagged, so from n = 5 they use 4 * sqrt(n), the typical rather than
worst-case error growth [3]. In bf16 the measured noise is lower still
and does not grow with n, so bf16 is also capped at 12 epsilon from
n = 10. The measurement is the ratio of the nonzero pivot that rounding
leaves in place of zero to the sum of the products subtracted to form it,
in epsilons, over 300 exactly singular matrices per size; the largest per
cell:

n	default backend, the sample below	default backend, a separate run	EXLA, a separate run
8	12.2	21.6	19.7
16	11.5	9.4	8.8
32	8.3	8.1	7.2
64	not in the sample	7.7	6.0

At 8x8 both rules allow 11.3, so the outliers there slip under either
rule; from 16x16 up the largest ratio stays under 12 everywhere. The
cap's cost is the tail: in four runs of 2,000 singular 16x16 matrices,
two per backend, 16 exceeded 12, about one in five hundred, with the
largest at 18.9; the sample below has none above 12 at 16x16, and an
earlier sample of 300 had 2. That is below the miss rate the uncapped
rules already have at 16x16, where f32, f64 and f16 miss 2, 6 and 6 of
300 in the table below because a computed LU can be the exact
factorization of a nearby nonsingular matrix. The cap's benefit is the
nonsingular matrices between 12 and 4 * sqrt(n) epsilon, which at 32x32
include the Hadamard matrix and the random well-conditioned families
below. f16 keeps the square root rule because its ratios are much larger:
in the separate run the largest per size is 22.6 at 16x16, 26.9 at 32x32
and 22.3 at 64x64, against allowances of 16, 22.6 and 32, which miss 5, 1
and 0 of the 300.

Torchx. invert/1 now factors instead of solving, so it needs the
Torchx lu NIF to factor singular matrices instead of throwing, which is
#1877.

Measurements

Below, cond is the 2-norm condition number of the matrix.

The next two tables and the breakdown after them come from one seeded
sample of 300 exactly singular and 900 nonsingular matrices per type at
sizes 2, 3, 4, 6, 8, 16 and 32, run once on main and once on the
branch. On main a matrix counts as caught or rejected when invert/1
returned NaN or raised, since the solve/2 it called raised on an
exactly zero pivot; the branch never raised. The 3x3 and 6x6 cells are
omitted from both tables, and 2x2 from the singular one, where every type
catches all 300.

Singular matrices caught. The branch catches nearly all of them at
every size and in every type; main caught most at 4x4, about a tenth
at 8x8 and almost none from 16x16 up. Random exactly singular integer
matrices, 300 per cell, main -> branch.

n	f32	f64	f16	bf16
4	202 -> 299	197 -> 299	276 -> 298	167 -> 296
8	25 -> 290	35 -> 297	41 -> 296	35 -> 299
16	9 -> 298	11 -> 294	6 -> 294	8 -> 300
32	2 -> 297	5 -> 297	8 -> 299	7 -> 300

The few singular matrices that still pass have a computed LU that is also
the exact factorization of a nearby nonsingular matrix, so no test on the
factors can flag them.

Nonsingular matrices wrongly rejected. f32 and f64 lose nothing, f16
loses far less than on main, and bf16 rejects more than main at
every size. 900 per cell: 300 each with small integer entries, standard
normal entries, and graded rows. main -> branch.

n	f32	f64	f16	bf16
2	0 -> 0	0 -> 0	281 -> 0	0 -> 2
4	0 -> 0	0 -> 0	309 -> 0	1 -> 15
8	7 -> 0	0 -> 0	300 -> 11	0 -> 73
16	287 -> 0	0 -> 0	300 -> 24	2 -> 226
32	300 -> 0	0 -> 0	300 -> 255	297 -> 670

The bf16 column is a trade, not a regression: of what bf16 now rejects,
main was inverting badly most at 32x32, half at 16x16 and fewer than
half at 8x8, so the caller gets NaN in place of wrong numbers, at the
cost of some inverses that were acceptable. The table shows how the
inverses main returned for those matrices had fared, in the max norm
against the f64 inverse of the same rounded matrix:

n	rejected	NaN or raise on main	more than 10% off	5% to 10% off	within 5%
8	73	0	30	11	32
16	226	2	115	45	64
32	670	255	299	96	20

At 32x32, 554 of the 670 rejected had failed or were more than 10% off
on main; at 16x16, 117 of 226, with 64 within 5%; at 8x8, 30 of 73,
with 32 within 5%. The 2 at 16x16 are one NaN and one ArgumentError,
where main's bf16 LU met an exactly zero pivot. The f16 counts at 2x2
and 4x4 are the small-entries bug from the Problem section disappearing;
f32 and f64 reject nothing at any size.

The same picture on 40 random normal bf16 matrices per size, from a
separate run measured the same way:

n	main: NaN, median error, worst	branch: NaN, median error, worst
8	0, 2.7%, 52%	6, 2.0%, 15%
16	0, 5.5%, 310%	11, 4.5%, 17%
32	0, 13%, 320%	28, 12%, 61%

The branch returns fewer inverses in bf16, and the worst it returns in
these samples is 61% off, against 320% on main.

Exhaustive checks. Every 2x2 with entries in -5..5 and every 3x3 with
entries in -1..1, in all four types, comes back NaN for exactly the
singular ones. On 20,000 random 4x4 with entries in -2..2, f32, f64 and
f16 are exact too; bf16 flags 122 of the 18,020 nonsingular matrices and
misses 3 of the 1,980 singular ones.

Well-conditioned and ill-conditioned matrices. Taking
cond * epsilon <= 0.1 as the range where a direct solver in the type can
be expected to succeed [6], the test rejects nothing in that range at
sizes 4 to 64 in f64, f32 and f16, and through 32x32 in bf16, on LAPACK's
dchkge families at cond 2, sqrt(0.1/epsilon) and 0.1/epsilon [4],
including a cond-2 matrix scaled to just above the smallest normal number,
and on thirteen named matrices from Higham's collection [5], Hadamard
included. The forward error of the accepted inverses grows as
cond * epsilon. Hilbert, Pascal, Frank and Vandermonde are rejected only
once cond * epsilon is far above 1. At 64x64, bf16 rejects 5 of 30
random matrices at cond 0.1/epsilon, the edge of the range, and nothing
else in range.

Entry magnitude. Scaling makes the result independent of the size of
the entries. With m * [[2, 1], [1, 3]] and m set to every power of
two representable in each type, the branch inverts it correctly wherever
the inverse is representable. main returned NaN for m below 2^-4 in
f16 and below 2^-60 in bf16 and f32. In f64 it returned NaN below
2^-508, where the determinant drops below the threshold, and raised
ArithmeticError above 2^510, where the determinant overflows.

Known limits. Two limits remain, and both belong to the backends
rather than to the test. On the default backend an f64 result
that overflows raises ArithmeticError, as every overflowing f64
operation does there. XLA flushes subnormal values to zero, so on EXLA
input below the normal range is zero. For f32 matrices whose rows exceed
about 2^126.5 in magnitude, an intermediate in XLA's triangular solve
is subnormal and is flushed, which drops the part of the inverse that
distinguishes the matrix from the nearest exactly singular one; the
relative error of the inverse is then about 1 / cond(A). That error is
largest for well-conditioned matrices, but their inverses at that
magnitude are subnormal themselves and would be flushed by XLA in any
case.

Behavior by backend

invert/1 now treats a singular matrix the same on all three backends:
NaN for that matrix, alone or in a batch, and no raise; Torchx still
raises for f16 and bf16 input, as before. Before this change, every
backend returned NaN only when the determinant fell below the threshold
and otherwise large finite numbers for a singular matrix, and the default
backend and Torchx raised for a batch containing a singular matrix.
After:

	default backend	EXLA	Torchx
invert/1, singular matrix alone or in a batch	NaN for that matrix	NaN for that matrix	NaN for that matrix (before: ArgumentError from Torchx's determinant check in solve/2)
invert/1, f16 or bf16 input	result in the input type (before: f32; #1876)	same	raises, before and after: LibTorch has no 16-bit LU on CPU

That is PyTorch's inv_ex model, a per-matrix status, with NaN as the
status since compiled code cannot raise. For comparison, PyTorch's inv
and solve raise only on an exactly zero pivot, and for the whole batch,
while inv_ex and solve_ex never raise. For
[[2, -1, -1, 0], [2, 0, -1, 1], [-1, 0, 1, 1], [1, 2, -2, -2]], which is
exactly singular with no exactly zero pivot, PyTorch returns entries of
order 1e7 with a success status. PyTorch rejects f16 and bf16 input to
these functions, NumPy rejects f16, and LAPACK has no 16-bit routines.

solve/2 is unchanged by this PR. On the default backend it raises only
when its triangular solve meets an exactly zero pivot; making it raise
for matrices that are singular to working precision, by the same test
invert/1 uses, is a follow-up PR.

Tests

nx/test/nx/lin_alg_test.exs adds invert/1 cases for:

• NaN in the input type for singular matrices alone and in a batch, with
  and without an exactly zero pivot, including samples of the 2x2 and 3x3
  domains checked exhaustively above and a determinant that rounds away
  from zero;
• well-conditioned matrices with small entries, graded matrices, and the
  cond-2 and cond-0.1/epsilon families at n <= 16, in every float type;
• bf16 diagonal and block-diagonal matrices at sizes where the tolerance
  is large;
• the Hadamard matrix at sizes up to 32, and the Wilkinson pivot-growth
  matrix [7], whose last pivot 2^(n - 1) overflows f16 at n = 17;
• signed integer input holding the type's minimum and rows below the
  normal range of f64;
• property tests on random small integer matrices: NaN for every singular
  one and an accurate inverse for every nonsingular one, alone and in a
  batch, with the matching solve/2 cases, which pass on main's
  solve/2 too because eliminating these small integers meets an exact
  zero pivot.

exla/ tests that the input type is kept, NaN per matrix in a batch, and
rows near the top of the f32 range. torchx/ tests that invert/1
returns NaN for singular matrices alone and in a batch, with and without
an exactly zero pivot.

Performance

invert/1 is about twice as fast from 8x8 up, because it no longer
computes a determinant in addition to the factorization. Default
backend, main -> branch:

case	invert/1
f32 2x2	1.84 -> 2.02 ms
f32 8x8	8.07 -> 4.30 ms
f32 32x32	83.4 -> 43.6 ms
f32 64x64	578 -> 304 ms
f64 64x64	630 -> 306 ms
f32, batch of 10 16x16	124 -> 70.9 ms

References

1. J. Demmel, Y. Hida, W. Kahan, X. S. Li, S. Mukherjee and E. J. Riedy,
"Error bounds from extra-precise iterative refinement", ACM Trans. Math.
Softw. 32(2), 2006. https://doi.org/10.1145/1141885.1141894
2. S. M. Rump and C.-P. Jeannerod, "Improved backward error bounds for LU
and Cholesky factorizations", SIAM J. Matrix Anal. Appl. 35(2), 2014.
https://doi.org/10.1137/130927231
3. N. J. Higham and T. Mary, "A new approach to probabilistic rounding
error analysis", SIAM J. Sci. Comput. 41(5), 2019.
https://doi.org/10.1137/18M1226312
4. G. W. Stewart, "The efficient generation of random orthogonal matrices
with an application to condition estimators", SIAM J. Numer. Anal.
17(3), 1980. https://doi.org/10.1137/0717034
5. N. J. Higham, "Algorithm 694: a collection of test matrices in MATLAB",
ACM Trans. Math. Softw. 17(3), 1991. https://doi.org/10.1145/114697.116805
6. N. J. Higham, Accuracy and Stability of Numerical Algorithms, 2nd ed.,
SIAM, 2002, chapter 7.
7. J. H. Wilkinson, "Error analysis of direct methods of matrix inversion",
J. ACM 8(3), 1961. https://doi.org/10.1145/321075.321076
8. N. J. Higham, "FORTRAN codes for estimating the one-norm of a real or
complex matrix, with applications to condition estimation", ACM Trans.
Math. Softw. 14(4), 1988. https://doi.org/10.1145/50063.214386

[polvalente]

is this correct? the name and the calculation seem disconnected

[heydtn]

It's correct, the transpose is what makes it hard to see and the naming could be better. Row i of
upper_magnitudes is column i of |U|, so summing the elementwise product
over row i gives the sum over j < i of |l[i][j]| * |u[j][i]|, the
magnitudes of the products subtracted from a[i][i] to form pivot i.

How about upper_column_magnitudes for the transposed U and
subtracted_magnitudes for the sum, with a comment stating that formula?

[polvalente]

Yeah, I still think calling it anything subtracted is a bit too hidden. magnitude_product might be a better name

[polvalente]

this seems overly specific to add to Nx.Type

[polvalente]

Suggested change

	If a singular matrix is passed, the result will silently fail. Matrices
	that are singular to working precision return NaN too, and so does one
	with a pivot that overflows the type. In a batch, only the singular
	matrices are NaN.
	If a singular matrix is passed, the result will silently fail;

I don't think these phrases add any new information:

• working precision is by definition how we can measure whether a matrix is singular
• overflows by definition silently fail
• batches are also defined as independently processed

[polvalente]

The PR is still a draft. Is it ready?

[heydtn]

The PR is still a draft. Is it ready?

Not yet - the numbers looked okay, but I wasn't satisfied. I did more extensive testing and found the source of some subtle drifting I noticed - there's several rounding vs truncation issues across the backends which cause this heuristic (as well as probably the original heuristic? or at least probably math somewhere) to be subtly almost correct. This heuristic expects rounding. Rounding follows IEEE 754, and it seems like rounding is the standard across the ecosystem. Literature seems to take this into account too? There were some places that were double rounding too, which causes issues.

The PR #1890 was a finding during this testing.

These issues impact other things as well I'm guessing, not just this heuristic, and subtly?

TLDR still working on things 😄