Move info!=0 logic into lax.linalg.tridiagonal lowering rule.

PiperOrigin-RevId: 726560133
jax-ml · Feb 13, 2025 · 011d297 · 011d297
1 parent 5ebb7eb
commit 011d297
Show file tree

Hide file tree

Showing 3 changed files with 346 additions and 254 deletions.
diff --git a/docs/jax.lax.rst b/docs/jax.lax.rst
@@ -238,19 +238,24 @@ Linear algebra operators (jax.lax.linalg)
   :toctree: _autosummary
 
     cholesky
+    cholesky_update
     eig
     eigh
     hessenberg
-    lu
     householder_product
+    lu
+    lu_pivots_to_permutation
     qdwh
     qr
     schur
     svd
+    SvdAlgorithm
+    symmetric_product
     triangular_solve
     tridiagonal
     tridiagonal_solve
 
+
 Argument classes
 ----------------