From c66f5dda60aa5df7b6aa2f09d3ce88c4249b6f34 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Tue, 15 Mar 2022 09:55:59 -0700
Subject: [PATCH] DOC: add missing linalg functionality to docs

---
 docs/jax.lax.rst       | 2 ++
 docs/jax.scipy.rst     | 4 ++++
 jax/_src/lax/linalg.py | 2 +-
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/docs/jax.lax.rst b/docs/jax.lax.rst
index 89f234dfe..4cbc040c2 100644
--- a/docs/jax.lax.rst
+++ b/docs/jax.lax.rst
@@ -207,8 +207,10 @@ Linear algebra operators (jax.lax.linalg)
     lu
     qdwh
     qr
+    schur
     svd
     triangular_solve
+    tridiagonal_solve
 
 Argument classes
 ----------------
diff --git a/docs/jax.scipy.rst b/docs/jax.scipy.rst
index 622dc30fe..96e557053 100644
--- a/docs/jax.scipy.rst
+++ b/docs/jax.scipy.rst
@@ -26,17 +26,21 @@ jax.scipy.linalg
    cholesky
    det
    eigh
+   eigh_tridiagonal
    expm
    expm_frechet
    inv
    lu
    lu_factor
    lu_solve
+   polar
+   polar_unitary
    qr
    schur
    sqrtm
    solve
    solve_triangular
+   sqrtm
    svd
    tril
    triu
diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
index 8be9e4ae3..e5ff62969 100644
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -1461,7 +1461,7 @@ xla.register_translation(tridiagonal_solve_p, xla.lower_fun(
 def tridiagonal_solve(dl, d, du, b):
   r"""Computes the solution of a tridiagonal linear system.
 
-  This function computes the solution of a tridiagonal linear system::
+  This function computes the solution of a tridiagonal linear system:
 
   .. math::
     A . X = B