tfp.substrates.jax.math.gram_schmidt

Implementation of the modified Gram-Schmidt orthonormalization algorithm.

tfp.substrates.jax.math.gram_schmidt(
    vectors, num_vectors=None
)

We assume here that the vectors are linearly independent. Zero vectors will be left unchanged, but will also consume an iteration against num_vectors.

From [1]: "MGS is numerically equivalent to Householder QR factorization applied to the matrix A augmented with a square matrix of zero elements on top."

Historical note, see [1]: "modified" Gram-Schmidt was derived by Laplace [2], for elimination and not as an orthogonalization algorithm. "Classical" Gram-Schmidt actually came later [2]. Classical Gram-Schmidt has a sometimes catastrophic loss of orthogonality for badly conditioned matrices, which is discussed further in [1].

References

[1] Bjorck, A. (1994). Numerics of gram-schmidt orthogonalization. Linear Algebra and Its Applications, 197, 297-316.

[2] P. S. Laplace, Thiorie Analytique des Probabilites. Premier Supple'ment, Mme. Courtier, Paris, 1816.

[3] E. Schmidt, über die Auflosung linearer Gleichungen mit unendlich vielen Unbekannten, Rend. Circ. Mat. Pulermo (1) 25:53-77 (1908).

vectors A Tensor of shape [..., d, n] of d-dim column vectors to orthonormalize.
num_vectors Optional, number of leading vectors of the result to make orthogonal. If unspecified, then num_vectors = n, implying that each vector except for the last will be used in sequence.

A Tensor of shape [..., d, n] corresponding to the orthonormalization.