Repository: google/jax
Pull request #21081 (2024-05-06 14:03:21)

Compute source maps when pretty-printing jaxprs.
This change is in preparation for adding support for emitting https://tc39.es/source-map/ source map information for jaxprs, so the relationship between a jaxpr and its Python code can be visualized using tooling for that purpose.

This change adds a new `source_map()` pretty printer document, which causes the pretty-printer to populate a source_map side output during pretty printing. The change also teaches the core jaxpr pretty printer to populate source map information on each equation.

@gspschmid


Patch:
diff --git a/.bazelrc b/.bazelrc
index 7abbecec4..b89f909f5 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -79,11 +79,6 @@ build:nvcc_clang --action_env=TF_CUDA_CLANG="1"
 build:nvcc_clang --action_env=TF_NVCC_CLANG="1"
 build:nvcc_clang --@local_config_cuda//:cuda_compiler=nvcc
 
-# Requires MSVC and LLVM to be installed
-build:win_clang --extra_toolchains=@local_config_cc//:cc-toolchain-x64_windows-clang-cl
-build:win_clang --extra_execution_platforms=//jax/tools/toolchains:x64_windows-clang-cl
-build:win_clang --compiler=clang-cl
-
 # Later Bazel flag values override earlier values.
 # TODO(jieying): remove enable_gpu and xla_python_enable_gpu from build:cuda
 # after the pluin is released.
diff --git a/.github/workflows/windows_ci_clang.yml b/.github/workflows/windows_ci_clang.yml
deleted file mode 100644
index b135a7abf..000000000
--- a/.github/workflows/windows_ci_clang.yml
+++ /dev/null
@@ -1,72 +0,0 @@
-name: CI - Windows CPU (trial)
-on:
-  schedule:
-  - cron: "0 12 * * *" # Daily at 12:00 UTC
-  workflow_dispatch: # allows triggering the workflow run manually
-
-env:
-  DISTUTILS_USE_SDK: 1
-  MSSdk: 1
-
-jobs:
-  win-wheels:
-    strategy:
-      fail-fast: true
-      matrix:
-        os: [windows-2019-32core]
-        arch: [AMD64]
-        pyver: ['3.9']
-    name: ${{ matrix.os }} CI build
-    runs-on: ${{ matrix.os }}
-
-    steps:
-    - name: Cancel Previous Runs
-      uses: styfle/cancel-workflow-action@85880fa0301c86cca9da44039ee3bb12d3bedbfa # ratchet: styfle/cancel-workflow-action@0.12.1
-      with:
-        access_token: ${{ github.token }}
-
-    - name: Install LLVM/Clang
-      run: choco install llvm --version=18.1.4 --yes
-
-    - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
-      with:
-        path: jax
-
-    - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
-      with:
-        repository: openxla/xla
-        path: xla
-
-    - uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # ratchet:actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.pyver }}
-        cache: 'pip'
-
-    - name: Build wheels
-      env:
-        # TODO(belitskiy): Remove the version once JAX is fully on 6.5.0
-        USE_BAZEL_VERSION: 6.5.0
-        BAZEL_VC: "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC"
-        JAXLIB_NIGHTLY: true  # Tag the wheels as dev versions
-      run: |
-        cd jax
-        python -m pip install -r build/test-requirements.txt
-        python -m pip uninstall -y matplotlib
-        python -m pip install --pre --upgrade numpy==2.0.0rc1 scipy==1.13.0
-        "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-        python.exe build\build.py `
-          ('--bazel_options=--override_repository=xla=${{ github.workspace }}\xla' -replace '\\','\\') `
-          --bazel_options=--color=yes `
-          --bazel_options=--config=win_clang
-
-    - name: Run tests
-      env:
-        JAX_ENABLE_CHECKS: true
-        JAX_SKIP_SLOW_TESTS: true
-        PY_COLORS: 1
-      run: |
-        cd jax
-        python -m pip install -e ${{ github.workspace }}\jax
-        python -m pip install --no-index --find-links ${{ github.workspace }}\jax\dist jaxlib
-        echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
-        pytest -n auto --tb=short tests examples
diff --git a/build/build.py b/build/build.py
index 463b71308..68dd47374 100755
--- a/build/build.py
+++ b/build/build.py
@@ -260,7 +260,7 @@ def write_bazelrc(*, python_bin_path, remote_build,
                   cuda_toolkit_path, cudnn_install_path,
                   cuda_version, cudnn_version, rocm_toolkit_path,
                   cpu, cuda_compute_capabilities,
-                  rocm_amdgpu_targets, target_cpu_features,
+                  rocm_amdgpu_targets, bazel_options, target_cpu_features,
                   wheel_cpu, enable_mkl_dnn, use_clang, clang_path,
                   clang_major_version, enable_cuda, enable_nccl, enable_rocm,
                   build_gpu_plugin, enable_mosaic_gpu):
@@ -279,11 +279,11 @@ def write_bazelrc(*, python_bin_path, remote_build,
       f.write(f'build --action_env CLANG_COMPILER_PATH="{clang_path}"\n')
       f.write(f'build --repo_env CC="{clang_path}"\n')
       f.write(f'build --repo_env BAZEL_COMPILER="{clang_path}"\n')
-      f.write('build --copt=-Wno-error=unused-command-line-argument\n')
-      if clang_major_version in (16, 17, 18):
+      bazel_options.append("--copt=-Wno-error=unused-command-line-argument\n")
+      if clang_major_version in (16, 17):
         # Necessary due to XLA's old version of upb. See:
         # https://github.com/openxla/xla/blob/c4277a076e249f5b97c8e45c8cb9d1f554089d76/.bazelrc#L505
-        f.write("build --copt=-Wno-gnu-offsetof-extensions\n")
+        bazel_options.append("--copt=-Wno-gnu-offsetof-extensions\n")
 
     if cuda_toolkit_path:
       tf_cuda_paths.append(cuda_toolkit_path)
@@ -316,6 +316,8 @@ def write_bazelrc(*, python_bin_path, remote_build,
     if cpu is not None:
       f.write(f"build --cpu={cpu}\n")
 
+    for o in bazel_options:
+      f.write(f"build {o}\n")
     if target_cpu_features == "release":
       if wheel_cpu == "x86_64":
         f.write("build --config=avx_windows\n" if is_windows()
@@ -527,8 +529,7 @@ def main():
   parser.add_argument(
       "--bazel_options",
       action="append", default=[],
-      help="Additional options to pass to the main Bazel command to be "
-           "executed, e.g. `run`.")
+      help="Additional options to pass to bazel.")
   parser.add_argument(
       "--output_path",
       default=os.path.join(cwd, "dist"),
@@ -646,6 +647,7 @@ def main():
       cpu=args.target_cpu,
       cuda_compute_capabilities=args.cuda_compute_capabilities,
       rocm_amdgpu_targets=args.rocm_amdgpu_targets,
+      bazel_options=args.bazel_options,
       target_cpu_features=args.target_cpu_features,
       wheel_cpu=wheel_cpu,
       enable_mkl_dnn=args.enable_mkl_dnn,
@@ -664,51 +666,41 @@ def main():
 
   print("\nBuilding XLA and installing it in the jaxlib source tree...")
 
-  command_base = (
-    bazel_path,
-    *args.bazel_startup_options,
-    "run",
-    "--verbose_failures=true",
-    *args.bazel_options,
-  )
   if not args.build_cuda_kernel_plugin and not args.build_cuda_pjrt_plugin:
-    build_cpu_wheel_command = [
-      *command_base,
-      "//jaxlib/tools:build_wheel", "--",
+    command = ([bazel_path] + args.bazel_startup_options +
+      ["run", "--verbose_failures=true"] +
+      ["//jaxlib/tools:build_wheel", "--",
       f"--output_path={output_path}",
       f"--jaxlib_git_hash={get_githash()}",
-      f"--cpu={wheel_cpu}"
-    ]
+      f"--cpu={wheel_cpu}"])
     if args.build_gpu_plugin:
-      build_cpu_wheel_command.append("--skip_gpu_kernels")
+      command.append("--skip_gpu_kernels")
     if args.editable:
-      build_cpu_wheel_command.append("--editable")
-    print(" ".join(build_cpu_wheel_command))
-    shell(build_cpu_wheel_command)
+      command += ["--editable"]
+    print(" ".join(command))
+    shell(command)
 
   if args.build_gpu_plugin or args.build_cuda_kernel_plugin:
-    build_cuda_kernels_command = [
-      *command_base,
-      "//jaxlib/tools:build_cuda_kernels_wheel", "--",
+    build_cuda_kernels_command = ([bazel_path] + args.bazel_startup_options +
+      ["run", "--verbose_failures=true"] +
+      ["//jaxlib/tools:build_cuda_kernels_wheel", "--",
       f"--output_path={output_path}",
       f"--jaxlib_git_hash={get_githash()}",
       f"--cpu={wheel_cpu}",
-      f"--cuda_version={args.gpu_plugin_cuda_version}"
-    ]
+      f"--cuda_version={args.gpu_plugin_cuda_version}"])
     if args.editable:
       build_cuda_kernels_command.append("--editable")
     print(" ".join(build_cuda_kernels_command))
     shell(build_cuda_kernels_command)
 
   if args.build_gpu_plugin or args.build_cuda_pjrt_plugin:
-    build_pjrt_plugin_command = [
-      *command_base,
-      "//jaxlib/tools:build_gpu_plugin_wheel", "--",
+    build_pjrt_plugin_command = ([bazel_path] + args.bazel_startup_options +
+      ["run", "--verbose_failures=true"] +
+      ["//jaxlib/tools:build_gpu_plugin_wheel", "--",
       f"--output_path={output_path}",
       f"--jaxlib_git_hash={get_githash()}",
       f"--cpu={wheel_cpu}",
-      f"--cuda_version={args.gpu_plugin_cuda_version}"
-    ]
+      f"--cuda_version={args.gpu_plugin_cuda_version}"])
     if args.editable:
       build_pjrt_plugin_command.append("--editable")
     print(" ".join(build_pjrt_plugin_command))
diff --git a/jax/_src/core.py b/jax/_src/core.py
index e845bbcaa..e34ddb8cf 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -149,51 +149,11 @@ class Jaxpr:
   def pretty_print(self, *, source_info=False, print_shapes=True,
                    custom_pp_eqn_rules=True, name_stack=False,
                    print_effects: bool = False, **kwargs):
-    context = JaxprPpContext()
-    settings = JaxprPpSettings(
-        source_info=source_info,
-        print_shapes=print_shapes,
-        custom_pp_eqn_rules=custom_pp_eqn_rules,
-        name_stack=name_stack,
-        print_effects=print_effects)
-
-    # Compute how many times each jaxpr is used.
-    names = defaultdict[Jaxpr, str](lambda: "jaxpr")
-    jaxpr_counts = Counter[Jaxpr]()
-    s = deque([self])
-    while s:
-      jaxpr = s.popleft()
-      jaxpr_counts[jaxpr] += 1
-      for eqn in jaxpr.eqns:
-        # TODO(slebedev): Come up with a more elaborate heuristic for name=.
-        name = eqn.params.get("name")
-        if name is None:
-          s.extend(jaxprs_in_params(eqn.params))
-          continue
-        name = name.strip("<>")  # <lambda> -> lambda
-        for subjaxpr in jaxprs_in_params(eqn.params):
-          s.append(subjaxpr)
-          names.setdefault(subjaxpr, name)
-
-    # Pull jaxprs occurring more than once to the top-level, making sure
-    # that their names are unique.
-    docs = []
-    name_counts = Counter[str]()
-    for jaxpr, c in jaxpr_counts.items():
-      if c == 1:
-        continue
-      name = names[jaxpr]
-      if (count := name_counts[name]) > 0:
-        name_counts[name] += 1
-        name += str(count)
-        name_counts[name] += 1
-      else:
-        name_counts[name] += 1
-      docs.append(pp_top_level_jaxpr(name, jaxpr, context, settings))
-      context.used_names.add(name)
-      context.top_level_jaxprs[jaxpr] = name
-    docs.append(pp_jaxpr(self, context, settings))
-    return pp.concat(docs).format(**kwargs)
+    doc = pp_toplevel_jaxpr(
+      self, source_info=source_info, print_shapes=print_shapes,
+      custom_pp_eqn_rules=custom_pp_eqn_rules, name_stack=name_stack,
+      print_effects=print_effects)
+    return doc.format(**kwargs)
 
   def _repr_pretty_(self, p, cycle):
     return p.text(self.pretty_print(use_color=True))
@@ -212,6 +172,7 @@ class Jaxpr:
     return jaxpr
 
 
+
 def join_effects(*effects: Effects) -> Effects:
   return set().union(*effects) if effects else no_effects
 
@@ -3164,6 +3125,55 @@ def _check_map(ctx_factory, prim, in_avals, params):
 
 # ------------------- Jaxpr printed representation -------------------
 
+def pp_toplevel_jaxpr(jaxpr_to_print, *, source_info=False, print_shapes=True,
+                      custom_pp_eqn_rules=True, name_stack=False,
+                      print_effects: bool = False) -> pp.Doc:
+    context = JaxprPpContext()
+    settings = JaxprPpSettings(
+        source_info=source_info,
+        print_shapes=print_shapes,
+        custom_pp_eqn_rules=custom_pp_eqn_rules,
+        name_stack=name_stack,
+        print_effects=print_effects)
+
+    # Compute how many times each jaxpr is used.
+    names = defaultdict[Jaxpr, str](lambda: "jaxpr")
+    jaxpr_counts = Counter[Jaxpr]()
+    s = deque([jaxpr_to_print])
+    while s:
+      jaxpr = s.popleft()
+      jaxpr_counts[jaxpr] += 1
+      for eqn in jaxpr.eqns:
+        # TODO(slebedev): Come up with a more elaborate heuristic for name=.
+        name = eqn.params.get("name")
+        if name is None:
+          s.extend(jaxprs_in_params(eqn.params))
+          continue
+        name = name.strip("<>")  # <lambda> -> lambda
+        for subjaxpr in jaxprs_in_params(eqn.params):
+          s.append(subjaxpr)
+          names.setdefault(subjaxpr, name)
+
+    # Pull jaxprs occurring more than once to the top-level, making sure
+    # that their names are unique.
+    docs = []
+    name_counts = Counter[str]()
+    for jaxpr, c in jaxpr_counts.items():
+      if c == 1:
+        continue
+      name = names[jaxpr]
+      if (count := name_counts[name]) > 0:
+        name_counts[name] += 1
+        name += str(count)
+        name_counts[name] += 1
+      else:
+        name_counts[name] += 1
+      docs.append(pp_top_level_jaxpr(name, jaxpr, context, settings))
+      context.used_names.add(name)
+      context.top_level_jaxprs[jaxpr] = name
+    docs.append(pp_jaxpr(jaxpr_to_print, context, settings))
+    return pp.concat(docs)
+
 
 class JaxprPpSettings(NamedTuple):
   print_shapes: bool = True
@@ -3253,7 +3263,9 @@ def pp_eqn(eqn: JaxprEqn, context: JaxprPpContext, settings: JaxprPpSettings
            ) -> pp.Doc:
   rule = (_pp_eqn if not settings.custom_pp_eqn_rules else
           pp_eqn_rules.get(eqn.primitive, _pp_eqn))
-  return rule(eqn, context, settings)  # type: ignore[operator]
+  doc = rule(eqn, context, settings)  # type: ignore[operator]
+  user_frame = source_info_util.user_frame(eqn.source_info)
+  return doc if user_frame is None else pp.source_map(doc, user_frame)
 
 def _pp_eqn(eqn, context, settings, params=None) -> pp.Doc:
   annotation = (source_info_util.summarize(eqn.source_info)
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 710387a09..2ce86faa1 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -19,6 +19,7 @@ from functools import partial
 import warnings
 
 import numpy as np
+import textwrap
 import operator
 from typing import Literal, NamedTuple, cast, overload
 
@@ -30,7 +31,7 @@ from jax._src.lax import lax as lax_internal
 from jax._src.lax import linalg as lax_linalg
 from jax._src.numpy import lax_numpy as jnp
 from jax._src.numpy import reductions, ufuncs
-from jax._src.numpy.util import promote_dtypes_inexact, check_arraylike
+from jax._src.numpy.util import implements, promote_dtypes_inexact, check_arraylike
 from jax._src.util import canonicalize_axis
 from jax._src.typing import ArrayLike, Array
 
@@ -63,67 +64,9 @@ def _H(x: ArrayLike) -> Array:
 def _symmetrize(x: Array) -> Array: return (x + _H(x)) / 2
 
 
+@implements(np.linalg.cholesky)
 @partial(jit, static_argnames=['upper'])
 def cholesky(a: ArrayLike, *, upper: bool = False) -> Array:
-  """Compute the Cholesky decomposition of a matrix.
-
-  JAX implementation of :func:`numpy.linalg.cholesky`.
-
-  The Cholesky decomposition of a matrix `A` is:
-
-  .. math::
-
-     A = U^HU
-
-  or
-
-  .. math::
-
-    A = LL^H
-
-  where `U` is an upper-triangular matrix and `L` is a lower-triangular matrix, and
-  :math:`X^H` is the Hermitian transpose of `X`.
-
-  Args:
-    a: input array, representing a (batched) positive-definite hermitian matrix.
-      Must have shape ``(..., N, N)``.
-    upper: if True, compute the upper Cholesky decomposition `L`. if False
-      (default), compute the lower Cholesky decomposition `U`.
-
-  Returns:
-    array of shape ``(..., N, N)`` representing the Cholesky decomposition
-    of the input. If the input is not Hermitian positive-definite, The result
-    will contain NaN entries.
-
-
-  See also:
-    - :func:`jax.scipy.linalg.cholesky`: SciPy-style Cholesky API
-    - :func:`jax.lax.linalg.cholesky`: XLA-style Cholesky API
-
-  Example:
-    A small real Hermitian positive-definite matrix:
-
-    >>> x = jnp.array([[2., 1.],
-    ...                [1., 2.]])
-
-    Lower Cholesky factorization:
-
-    >>> jnp.linalg.cholesky(x)
-    Array([[1.4142135 , 0.        ],
-           [0.70710677, 1.2247449 ]], dtype=float32)
-
-    Upper Cholesky factorization:
-
-    >>> jnp.linalg.cholesky(x, upper=True)
-    Array([[1.4142135 , 0.70710677],
-           [0.        , 1.2247449 ]], dtype=float32)
-
-    Reconstructing ``x`` from its factorization:
-
-    >>> L = jnp.linalg.cholesky(x)
-    >>> jnp.allclose(x, L @ L.T)
-    Array(True, dtype=bool)
-  """
   check_arraylike("jnp.linalg.cholesky", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   L = lax_linalg.cholesky(a)
@@ -187,6 +130,7 @@ def svd(
   ...
 
 
+@implements(np.linalg.svd)
 @partial(
     jit,
     static_argnames=(
@@ -203,75 +147,6 @@ def svd(
     hermitian: bool = False,
     subset_by_index: tuple[int, int] | None = None,
 ) -> Array | SVDResult:
-  r"""Compute the singular value decomposition.
-
-  JAX implementation of :func:`numpy.linalg.svd`, implemented in terms of
-  :func:`jax.lax.linalg.svd`.
-
-  The SVD of a matrix `A` is given by
-
-  .. math::
-
-     A = U\Sigma V^H
-
-  - :math:`U` contains the left singular vectors and satisfies :math:`U^HU=I`
-  - :math:`V` contains the right singular vectors and satisfies :math:`V^HV=I`
-  - :math:`\Sigma` is a diagonal matrix of singular values.
-
-  Args:
-    a: input array, of shape ``(..., N, M)``
-    full_matrices: if True (default) compute the full matrices; i.e. ``u`` and ``vh`` have
-      shape ``(..., N, N)`` and ``(..., M, M)``. If False, then the shapes are
-      ``(..., N, K)`` and ``(..., K, M)`` with ``K = min(N, M)``.
-    compute_uv: if True (default), return the full SVD ``(u, s, vh)``. If False then return
-      only the singular values ``s``.
-    hermitian: if True, assume the matrix is hermitian, which allows for a more efficient
-      implementation (default=False)
-    subset_by_index: (TPU-only) Optional 2-tuple [start, end] indicating the range of
-      indices of singular values to compute. For example, if ``[n-2, n]`` then
-      ``svd`` computes the two largest singular values and their singular vectors.
-      Only compatible with ``full_matrices=False``.
-
-  Returns:
-    A tuple of arrays ``(u, s, vh)`` if ``compute_uv`` is True, otherwise the array ``s``.
-
-    - ``u``: left singular vectors of shape ``(..., N, N)`` if ``full_matrices`` is True
-      or ``(..., N, K)`` otherwise.
-    - ``s``: singular values of shape ``(..., K)``
-    - ``vh``: conjugate-transposed right singular vectors of shape ``(..., M, M)``
-      if ``full_matrices`` is True or ``(..., K, M)`` otherwise.
-
-    where ``K = min(N, M)``.
-
-  See also:
-    - :func:`jax.scipy.linalg.svd`: SciPy-style SVD API
-    - :func:`jax.lax.linalg.svd`: XLA-style SVD API
-
-  Example:
-    Consider the SVD of a small real-valued array:
-
-    >>> x = jnp.array([[1., 2., 3.],
-    ...                [6., 5., 4.]])
-    >>> u, s, vt = jnp.linalg.svd(x, full_matrices=False)
-    >>> s  # doctest: +SKIP
-    Array([9.361919 , 1.8315067], dtype=float32)
-
-    The singular vectors are in the columns of ``u`` and ``v = vt.T``. These vectors are
-    orthonormal, which can be demonstrated by comparing the matrix product with the
-    identity matrix:
-
-    >>> jnp.allclose(u.T @ u, jnp.eye(2), atol=1E-5)
-    Array(True, dtype=bool)
-    >>> v = vt.T
-    >>> jnp.allclose(v.T @ v, jnp.eye(2), atol=1E-5)
-    Array(True, dtype=bool)
-
-    Given the SVD, ``x`` can be reconstructed via matrix multiplication:
-
-    >>> x_reconstructed = u @ jnp.diag(s) @ vt
-    >>> jnp.allclose(x_reconstructed, x)
-    Array(True, dtype=bool)
-  """
   check_arraylike("jnp.linalg.svd", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   if hermitian:
@@ -307,51 +182,9 @@ def svd(
     )
 
 
+@implements(np.linalg.matrix_power)
 @partial(jit, static_argnames=('n',))
 def matrix_power(a: ArrayLike, n: int) -> Array:
-  """Raise a square matrix to an integer power.
-
-  JAX implementation of :func:`numpy.linalg.matrix_power`, implemented via
-  repeated squarings.
-
-  Args:
-    a: array of shape ``(..., M, M)`` to be raised to the power `n`.
-    n: the integer exponent to which the matrix should be raised.
-
-  Returns:
-    Array of shape ``(..., M, M)`` containing the matrix power of a to the n.
-
-  Examples:
-    >>> a = jnp.array([[1., 2.],
-    ...                [3., 4.]])
-    >>> jnp.linalg.matrix_power(a, 3)
-    Array([[ 37.,  54.],
-           [ 81., 118.]], dtype=float32)
-    >>> a @ a @ a  # equivalent evaluated directly
-    Array([[ 37.,  54.],
-           [ 81., 118.]], dtype=float32)
-
-    This also supports zero powers:
-
-    >>> jnp.linalg.matrix_power(a, 0)
-    Array([[1., 0.],
-           [0., 1.]], dtype=float32)
-
-    and also supports negative powers:
-
-    >>> with jnp.printoptions(precision=3):
-    ...   jnp.linalg.matrix_power(a, -2)
-    Array([[ 5.5 , -2.5 ],
-           [-3.75,  1.75]], dtype=float32)
-
-    Negative powers are equivalent to matmul of the inverse:
-
-    >>> inv_a = jnp.linalg.inv(a)
-    >>> with jnp.printoptions(precision=3):
-    ...   inv_a @ inv_a
-    Array([[ 5.5 , -2.5 ],
-           [-3.75,  1.75]], dtype=float32)
-  """
   check_arraylike("jnp.linalg.matrix_power", a)
   arr, = promote_dtypes_inexact(jnp.asarray(a))
 
@@ -388,40 +221,9 @@ def matrix_power(a: ArrayLike, n: int) -> Array:
   return result
 
 
+@implements(np.linalg.matrix_rank)
 @jit
 def matrix_rank(M: ArrayLike, tol: ArrayLike | None = None) -> Array:
-  """Compute the rank of a matrix.
-
-  JAX implementation of :func:`numpy.linalg.matrix_rank`.
-
-  The rank is calculated via the Singular Value Decomposition (SVD), and determined
-  by the number of singular values greater than the specified tolerance.
-
-  Args:
-    a: array of shape ``(..., M, N)`` whose rank is to be computed.
-    tol: optional array of shape ``(...)`` specifying the tolerance. Singular values
-      smaller than `tol` are considered to be zero. If ``tol`` is None (the default),
-      a reasonable default is chosen based the floating point precision of the input.
-
-  Returns:
-    array of shape ``a.shape[-2]`` giving the matrix rank.
-
-  Notes:
-    The rank calculation may be inaccurate for matrices with very small singular
-    values or those that are numerically ill-conditioned. Consider adjusting the
-    ``tol`` parameter or using a more specialized rank computation method in such cases.
-
-  Examples:
-    >>> a = jnp.array([[1, 2],
-    ...                [3, 4]])
-    >>> jnp.linalg.matrix_rank(a)
-    Array(2, dtype=int32)
-
-    >>> b = jnp.array([[1, 0],  # Rank-deficient matrix
-    ...                [0, 0]])
-    >>> jnp.linalg.matrix_rank(b)
-    Array(1, dtype=int32)
-  """
   check_arraylike("jnp.linalg.matrix_rank", M)
   M, = promote_dtypes_inexact(jnp.asarray(M))
   if M.ndim < 2:
@@ -476,44 +278,23 @@ def _slogdet_qr(a: Array) -> tuple[Array, Array]:
   sign_taus = reductions.prod(jnp.where(taus[..., :(n-1)] != 0, -1, 1), axis=-1).astype(sign_diag.dtype)
   return sign_diag * sign_taus, log_abs_det
 
-
+@implements(
+    np.linalg.slogdet,
+    extra_params=textwrap.dedent("""
+      method: string, optional
+        One of ``lu`` or ``qr``, specifying whether the determinant should be
+        computed using an LU decomposition or a QR decomposition. Defaults to
+        LU decomposition if ``None``.
+    """))
 @partial(jit, static_argnames=('method',))
 def slogdet(a: ArrayLike, *, method: str | None = None) -> SlogdetResult:
-  """
-  Computes the sign and (natural) logarithm of the determinant of an array.
-
-  JAX implementation of :func:`numpy.linalg.slotdet`.
-
-  Args:
-    a: array of shape ``(..., M, M)`` for which to compute the sign and log determinant.
-    method: the method to use for determinant computation. Options are
-
-      - ``'lu'`` (default): use the LU decomposition.
-      - ``'qr'``: use the QR decomposition.
-
-  Returns:
-    A tuple of arrays ``(sign, logabsdet)``, each of shape ``a.shape[:-2]``
-
-    - ``sign`` is the sign of the determinant.
-    - ``logabsdet`` is the natural log of the determinant's absolute value.
-
-  See also:
-    :func:`jax.numpy.linalg.det`: direct computation of determinant
-
-  Examples:
-    >>> a = jnp.array([[1, 2],
-    ...                [3, 4]])
-    >>> sign, logabsdet = jnp.linalg.slogdet(a)
-    >>> sign  # -1 indicates negative determinant
-    Array(-1., dtype=float32)
-    >>> jnp.exp(logabsdet)  # Absolute value of determinant
-    Array(2., dtype=float32)
-  """
   check_arraylike("jnp.linalg.slogdet", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   a_shape = jnp.shape(a)
   if len(a_shape) < 2 or a_shape[-1] != a_shape[-2]:
-    raise ValueError("Argument to slogdet() must have shape [..., n, n], got {a_shape}")
+    msg = "Argument to slogdet() must have shape [..., n, n], got {}"
+    raise ValueError(msg.format(a_shape))
+
   if method is None or method == "lu":
     return SlogdetResult(*_slogdet_lu(a))
   elif method == "qr":
@@ -643,28 +424,9 @@ def _det_3x3(a: Array) -> Array:
 
 
 @custom_jvp
+@implements(np.linalg.det)
 @jit
 def det(a: ArrayLike) -> Array:
-  """
-  Computes the determinant of an array.
-
-  JAX implementation of :func:`numpy.linalg.det`.
-
-  Args:
-    a: array of shape ``(..., M, M)`` for which to compute the determinant.
-
-  Returns:
-    An array of determinants of shape ``a.shape[:-2]``.
-
-  See also:
-    :func:`jax.scipy.linalg.det`: Scipy-style API for determinant.
-
-  Examples:
-    >>> a = jnp.array([[1, 2],
-    ...                [3, 4]])
-    >>> jnp.linalg.det(a)
-    Array(-2., dtype=float32)
-  """
   check_arraylike("jnp.linalg.det", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   a_shape = jnp.shape(a)
@@ -688,127 +450,34 @@ def _det_jvp(primals, tangents):
   return y, jnp.trace(z, axis1=-1, axis2=-2)
 
 
-def eig(a: ArrayLike) -> tuple[Array, Array]:
-  """
-  Computes the eigenvalues and eigenvectors of a square array.
-
-  JAX implementation of :func:`numpy.linalg.eig`.
-
-  Args:
-    a: array of shape ``(..., M, M)`` for which to compute the eigenvalues and vectors.
+@implements(np.linalg.eig, lax_description="""
+This differs from :func:`numpy.linalg.eig` in that the return type of
+:func:`jax.numpy.linalg.eig` is always ``complex64`` for 32-bit input,
+and ``complex128`` for 64-bit input.
 
-  Returns:
-    A tuple ``(eigenvalues, eigenvectors)`` with
-
-    - ``eigenvalues``: an array of shape ``(..., M)`` containing the eigenvalues.
-    - ``eigenvectors``: an array of shape ``(..., M, M)``, where column ``v[:, i]`` is the
-      eigenvector corresponding to the eigenvalue ``w[i]``.
-
-  Notes:
-    - This differs from :func:`numpy.linalg.eig` in that the return type of
-      :func:`jax.numpy.linalg.eig` is always complex64 for 32-bit input, and complex128
-      for 64-bit input.
-    - At present, non-symmetric eigendecomposition is only implemented on the CPU backend.
-
-  See also:
-    - :func:`jax.numpy.linalg.eigh`: eigenvectors and eigenvalues of a Hermitian matrix.
-    - :func:`jax.numpy.linalg.eigvals`: compute eigenvalues only.
-
-  Examples:
-    >>> a = jnp.array([[1., 2.],
-    ...                [2., 1.]])
-    >>> w, v = jnp.linalg.eig(a)
-    >>> with jax.numpy.printoptions(precision=4):
-    ...   w
-    Array([ 3.+0.j, -1.+0.j], dtype=complex64)
-    >>> v
-    Array([[ 0.70710677+0.j, -0.70710677+0.j],
-           [ 0.70710677+0.j,  0.70710677+0.j]], dtype=complex64)
-  """
+At present, non-symmetric eigendecomposition is only implemented on the CPU
+backend. However eigendecomposition for symmetric/Hermitian matrices is
+implemented more widely (see :func:`jax.numpy.linalg.eigh`).
+""")
+def eig(a: ArrayLike) -> tuple[Array, Array]:
   check_arraylike("jnp.linalg.eig", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   w, v = lax_linalg.eig(a, compute_left_eigenvectors=False)
   return w, v
 
 
+@implements(np.linalg.eigvals)
 @jit
 def eigvals(a: ArrayLike) -> Array:
-  """
-  Computes the eigenvalues of a general matrix.
-
-  JAX implementation of :func:`numpy.linalg.eigvals`.
-
-  Args:
-    a: array of shape ``(..., M, M)`` for which to compute the eigenvalues.
-
-  Returns:
-    An array of shape ``(..., M)`` containing the eigenvalues.
-
-  See also:
-    - :func:`jax.numpy.linalg.eig`: computes eigenvalues eigenvectors of a general matrix.
-    - :func:`jax.numpy.linalg.eigh`: computes eigenvalues eigenvectors of a Hermitian matrix.
-
-  Notes:
-    - This differs from :func:`numpy.linalg.eigvals` in that the return type of
-      :func:`jax.numpy.linalg.eigvals` is always complex64 for 32-bit input, and
-      complex128 for 64-bit input.
-    - At present, non-symmetric eigendecomposition is only implemented on the CPU backend.
-
-  Examples:
-    >>> a = jnp.array([[1., 2.],
-    ...                [2., 1.]])
-    >>> w = jnp.linalg.eigvals(a)
-    >>> with jnp.printoptions(precision=2):
-    ...  w
-    Array([ 3.+0.j, -1.+0.j], dtype=complex64)
-  """
   check_arraylike("jnp.linalg.eigvals", a)
-  a, = promote_dtypes_inexact(jnp.asarray(a))
   return lax_linalg.eig(a, compute_left_eigenvectors=False,
                         compute_right_eigenvectors=False)[0]
 
 
+@implements(np.linalg.eigh)
 @partial(jit, static_argnames=('UPLO', 'symmetrize_input'))
 def eigh(a: ArrayLike, UPLO: str | None = None,
          symmetrize_input: bool = True) -> EighResult:
-  """
-  Computes the eigenvalues and eigenvectors of a Hermitian matrix.
-
-  JAX implementation of :func:`numpy.linalg.eigh`.
-
-  Args:
-    a: array of shape ``(..., M, M)``, containing the Hermitian (if complex)
-      or symmetric (if real) matrix.
-    UPLO: specifies whether the calculation isdone with the lower triangular
-      part of ``a`` (``'L'``, default) or the upper triangular part (``'U'``).
-    symmetrize_input: if True (default) then input is symmetrized, which leads
-      to better behavior under automatic differentiation.
-
-  Returns:
-    A namedtuple ``(eigenvalues, eigenvectors)`` where
-
-    - ``eigenvalues``: an array of shape ``(..., M)`` containing the eigenvalues,
-      sorted in ascending order.
-    - ``eigenvectors``: an array of shape ``(..., M, M)``, where column ``v[:, i]`` is the
-      normalized eigenvector corresponding to the eigenvalue ``w[i]``.
-
-  See also:
-    - :func:`jax.numpy.linalg.eig`: general eigenvalue decomposition.
-    - :func:`jax.numpy.linalg.eigvalsh`: compute eigenvalues only.
-    - :func:`jax.scipy.linalg.eigh`: SciPy API for Hermitian eigendecomposition.
-    - :func:`jax.lax.linalg.eigh`: XLA API for Hermitian eigendecomposition.
-
-  Examples:
-    >>> a = jnp.array([[1, -2j],
-    ...                [2j, 1]])
-    >>> w, v = jnp.linalg.eigh(a)
-    >>> w
-    Array([-1.,  3.], dtype=float32)
-    >>> with jnp.printoptions(precision=3):
-    ...   v
-    Array([[-0.707+0.j   , -0.707+0.j   ],
-           [ 0.   +0.707j,  0.   -0.707j]], dtype=complex64)
-  """
   check_arraylike("jnp.linalg.eigh", a)
   if UPLO is None or UPLO == "L":
     lower = True
@@ -823,89 +492,27 @@ def eigh(a: ArrayLike, UPLO: str | None = None,
   return EighResult(w, v)
 
 
+@implements(np.linalg.eigvalsh)
 @partial(jit, static_argnames=('UPLO',))
 def eigvalsh(a: ArrayLike, UPLO: str | None = 'L') -> Array:
-  """
-  Computes the eigenvalues of a Hermitian matrix.
-
-  JAX implementation of :func:`numpy.linalg.eigvalsh`.
-
-  Args:
-    a: array of shape ``(..., M, M)``, containing the Hermitian (if complex)
-      or symmetric (if real) matrix.
-    UPLO: specifies whether the calculation is done with the lower triangular
-      part of ``a`` (``'L'``, default) or the upper triangular part (``'U'``).
-
-  Returns:
-    An array of shape ``(..., M)`` containing the eigenvalues, sorted in
-    ascending order.
-
-  See also:
-    - :func:`jax.numpy.linalg.eig`: general eigenvalue decomposition.
-    - :func:`jax.numpy.linalg.eigh`: computes eigenvalues and eigenvectors of a
-      Hermitian matrix.
-
-  Examples:
-    >>> a = jnp.array([[1, -2j],
-    ...                [2j, 1]])
-    >>> w = jnp.linalg.eigvalsh(a)
-    >>> w
-    Array([-1.,  3.], dtype=float32)
-  """
   check_arraylike("jnp.linalg.eigvalsh", a)
-  a, = promote_dtypes_inexact(jnp.asarray(a))
   w, _ = eigh(a, UPLO)
   return w
 
 
 @partial(custom_jvp, nondiff_argnums=(1, 2))
+@implements(np.linalg.pinv, lax_description=textwrap.dedent("""\
+    It differs only in default value of `rcond`. In `numpy.linalg.pinv`, the
+    default `rcond` is `1e-15`. Here the default is
+    `10. * max(num_rows, num_cols) * jnp.finfo(dtype).eps`.
+    """))
 @partial(jit, static_argnames=('hermitian',))
 def pinv(a: ArrayLike, rcond: ArrayLike | None = None,
          hermitian: bool = False) -> Array:
-  """Compute the (Moore-Penrose) pseudo-inverse of a matrix.
-
-  JAX implementation of :func:`numpy.linalg.pinv`.
-
-  Args:
-    a: array of shape ``(..., M, N)`` containing matrices to pseudo-invert.
-    rcond: float or array_like of shape ``a.shape[:-2]``. Specifies the cutoff
-      for small singular values.of shape ``(...,)``.
-      Cutoff for small singular values; singular values smaller than this value
-      are treated as zero. The default is determined based on the floating point
-      precision of the dtype.
-    hermitian: if True, then the input is assumed to be Hermitian, and a more
-      efficient algorithm is used (default: False)
-
-  Returns:
-    An array of shape ``(..., N, M)`` containing the pseudo-inverse of ``a``.
-
-  See also:
-    - :func:`jax.numpy.linalg.inv`: multiplicative inverse of a square matrix.
-
-  Notes:
-    :func:`jax.numpy.linalg.prng` differs from :func:`numpy.linalg.prng` in the
-    default value of `rcond``: in NumPy, the default  is `1e-15`. In JAX, the
-    default is ``10. * max(num_rows, num_cols) * jnp.finfo(dtype).eps``.
-
-  Examples:
-    >>> a = jnp.array([[1, 2],
-    ...                [3, 4],
-    ...                [5, 6]])
-    >>> a_pinv = jnp.linalg.pinv(a)
-    >>> a_pinv  # doctest: +SKIP
-    Array([[-1.333332  , -0.33333257,  0.6666657 ],
-           [ 1.0833322 ,  0.33333272, -0.41666582]], dtype=float32)
-
-    The pseudo-inverse operates as a multiplicative inverse so long as the
-    output is not rank-deficient:
-
-    >>> jnp.allclose(a_pinv @ a, jnp.eye(2), atol=1E-4)
-    Array(True, dtype=bool)
-  """
   # Uses same algorithm as
   # https://github.com/numpy/numpy/blob/v1.17.0/numpy/linalg/linalg.py#L1890-L1979
   check_arraylike("jnp.linalg.pinv", a)
-  arr, = promote_dtypes_inexact(jnp.asarray(a))
+  arr = jnp.asarray(a)
   m, n = arr.shape[-2:]
   if m == 0 or n == 0:
     return jnp.empty(arr.shape[:-2] + (n, m), arr.dtype)
@@ -954,57 +561,9 @@ def _pinv_jvp(rcond, hermitian, primals, tangents):
   return p, p_dot
 
 
+@implements(np.linalg.inv)
 @jit
 def inv(a: ArrayLike) -> Array:
-  """Return the inverse of a square matrix
-
-  JAX implementation of :func:`numpy.linalg.inv`.
-
-  Args:
-    a: array of shape ``(..., N, N)`` specifying square array(s) to be inverted.
-
-  Returns:
-    Array of shape ``(..., N, N)`` containing the inverse of the input.
-
-  Notes:
-    In most cases, explicitly computing the inverse of a matrix is ill-advised. For
-    example, to compute ``x = inv(A) @ b``, it is more performant and numerically
-    precise to use a direct solve, such as :func:`jax.scipy.linalg.solve`.
-
-  See Also:
-    - :func:`jax.scipy.linalg.inv`: SciPy-style API for matrix inverse
-    - :func:`jax.numpy.linalg.solve`: direct linear solver
-
-  Example:
-    Compute the inverse of a 3x3 matrix
-
-    >>> a = jnp.array([[1., 2., 3.],
-    ...                [2., 4., 2.],
-    ...                [3., 2., 1.]])
-    >>> a_inv = jnp.linalg.inv(a)
-    >>> a_inv  # doctest: +SKIP
-    Array([[ 0.        , -0.25      ,  0.5       ],
-           [-0.25      ,  0.5       , -0.25000003],
-           [ 0.5       , -0.25      ,  0.        ]], dtype=float32)
-
-    Check that multiplying with the inverse gives the identity:
-
-    >>> jnp.allclose(a @ a_inv, jnp.eye(3), atol=1E-5)
-    Array(True, dtype=bool)
-
-    Multiply the inverse by a vector ``b``, to find a solution to ``a @ x = b``
-
-    >>> b = jnp.array([1., 4., 2.])
-    >>> a_inv @ b
-    Array([ 0.  ,  1.25, -0.5 ], dtype=float32)
-
-    Note, however, that explicitly computing the inverse in such a case can lead
-    to poor performance and loss of precision as the size of the problem grows.
-    Instead, you should use a direct solver like :func:`jax.numpy.linalg.solve`:
-
-    >>> jnp.linalg.solve(a, b)
-     Array([ 0.  ,  1.25, -0.5 ], dtype=float32)
-  """
   check_arraylike("jnp.linalg.inv", a)
   arr = jnp.asarray(a)
   if arr.ndim < 2 or arr.shape[-1] != arr.shape[-2]:
@@ -1014,74 +573,11 @@ def inv(a: ArrayLike) -> Array:
     arr, lax.broadcast(jnp.eye(arr.shape[-1], dtype=arr.dtype), arr.shape[:-2]))
 
 
+@implements(np.linalg.norm)
 @partial(jit, static_argnames=('ord', 'axis', 'keepdims'))
 def norm(x: ArrayLike, ord: int | str | None = None,
          axis: None | tuple[int, ...] | int = None,
          keepdims: bool = False) -> Array:
-  """Compute the norm of a matrix or vector.
-
-  JAX implementation of :func:`numpy.linalg.norm`.
-
-  Args:
-    x: N-dimensional array for which the norm will be computed.
-    ord: specify the kind of norm to take. Default is Frobenius norm for matrices,
-      and the 2-norm for vectors. For other options, see Notes below.
-    axis: integer or sequence of integers specifying the axes over which the norm
-      will be computed. Defaults to all axes of ``x``.
-    keepdims: if True, the output array will have the same number of dimensions as
-      the input, with the size of reduced axes replaced by ``1`` (default: False).
-
-  Returns:
-    array containing the specified norm of x.
-
-  Notes:
-    The flavor of norm computed depends on the value of ``ord`` and the number of
-    axes being reduced.
-
-    For **vector norms** (i.e. a single axis reduction):
-
-    - ``ord=None`` (default) computes the 2-norm
-    - ``ord=inf`` computes ``max(abs(x))``
-    - ``ord=-inf`` computes min(abs(x))``
-    - ``ord=0`` computes ``sum(x!=0)``
-    - for other numerical values, computes ``sum(abs(x) ** ord)**(1/ord)``
-
-    For **matrix norms** (i.e. two axes reductions):
-
-    - ``ord='fro'`` or ``ord=None`` (default) computes the Frobenius norm
-    - ``ord='nuc'`` computes the nuclear norm, or the sum of the singular values
-    - ``ord=1`` computes ``max(abs(x).sum(0))``
-    - ``ord=-1`` computes ``min(abs(x).sum(0))``
-    - ``ord=2`` computes the 2-norm, i.e. the largest singular value
-    - ``ord=-2`` computes the smallest singular value
-
-  Examples:
-    Vector norms:
-
-    >>> x = jnp.array([3., 4., 12.])
-    >>> jnp.linalg.norm(x)
-    Array(13., dtype=float32)
-    >>> jnp.linalg.norm(x, ord=1)
-    Array(19., dtype=float32)
-    >>> jnp.linalg.norm(x, ord=0)
-    Array(3., dtype=float32)
-
-    Matrix norms:
-
-    >>> x = jnp.array([[1., 2., 3.],
-    ...                [4., 5., 7.]])
-    >>> jnp.linalg.norm(x)  # Frobenius norm
-    Array(10.198039, dtype=float32)
-    >>> jnp.linalg.norm(x, ord='nuc')  # nuclear norm
-    Array(10.762535, dtype=float32)
-    >>> jnp.linalg.norm(x, ord=1)  # 1-norm
-    Array(10., dtype=float32)
-
-    Batched vector norm:
-
-    >>> jnp.linalg.norm(x, axis=1)
-    Array([3.7416575, 9.486833 ], dtype=float32)
-  """
   check_arraylike("jnp.linalg.norm", x)
   x, = promote_dtypes_inexact(jnp.asarray(x))
   x_shape = jnp.shape(x)
@@ -1179,72 +675,9 @@ def qr(a: ArrayLike, mode: Literal["r"]) -> Array: ...
 @overload
 def qr(a: ArrayLike, mode: str = "reduced") -> Array | QRResult: ...
 
+@implements(np.linalg.qr)
 @partial(jit, static_argnames=('mode',))
 def qr(a: ArrayLike, mode: str = "reduced") -> Array | QRResult:
-  """Compute the QR decomposition of an array
-
-  JAX implementation of :func:`numpy.linalg.qr`.
-
-  The QR decomposition of a matrix `A` is given by
-
-  .. math::
-
-     A = QR
-
-  Where `Q` is a unitary matrix (i.e. :math:`Q^HQ=I`) and `R` is an upper-triangular
-  matrix.
-
-  Args:
-    a: array of shape (..., M, N)
-    mode: Computational mode. Supported values are:
-
-      - ``"reduced"`` (default): return `Q` of shape ``(..., M, K)`` and `R` of shape
-        ``(..., K, N)``, where ``K = min(M, N)``.
-      - ``"complete"``: return `Q` of shape ``(..., M, M)`` and `R` of shape ``(..., M, N)``.
-      - ``"raw"``: return lapack-internal representations of shape ``(..., M, N)`` and ``(..., K)``.
-      - ``"r"``: return `R` only.
-
-  Returns:
-    A tuple ``(Q, R)`` (if ``mode`` is not ``"r"``) otherwise an array ``R``,
-    where:
-
-    - ``Q`` is an orthogonal matrix of shape ``(..., M, K)`` (if ``mode`` is ``"reduced"``)
-      or ``(..., M, M)`` (if ``mode`` is ``"complete"``).
-    - ``R`` is an upper-triangular matrix of shape ``(..., M, N)`` (if ``mode`` is
-      ``"r"`` or ``"complete"``) or ``(..., K, N)`` (if ``mode`` is ``"reduced"``)
-
-    with ``K = min(M, N)``.
-
-  See also:
-    - :func:`jax.scipy.linalg.qr`: SciPy-style QR decomposition API
-    - :func:`jax.lax.linalg.qr`: XLA-style QR decompositon API
-
-  Examples:
-    Compute the QR decomposition of a matrix:
-
-    >>> a = jnp.array([[1., 2., 3., 4.],
-    ...                [5., 4., 2., 1.],
-    ...                [6., 3., 1., 5.]])
-    >>> Q, R = jnp.linalg.qr(a)
-    >>> Q  # doctest: +SKIP
-    Array([[-0.12700021, -0.7581426 , -0.6396022 ],
-           [-0.63500065, -0.43322435,  0.63960224],
-           [-0.7620008 ,  0.48737738, -0.42640156]], dtype=float32)
-    >>> R  # doctest: +SKIP
-    Array([[-7.8740077, -5.080005 , -2.4130025, -4.953006 ],
-           [ 0.       , -1.7870499, -2.6534991, -1.028908 ],
-           [ 0.       ,  0.       , -1.0660033, -4.050814 ]], dtype=float32)
-
-    Check that ``Q`` is orthonormal:
-
-    >>> jnp.allclose(Q.T @ Q, jnp.eye(3), atol=1E-5)
-    Array(True, dtype=bool)
-
-    Reconstruct the input:
-
-    >>> jnp.allclose(Q @ R, a)
-    Array(True, dtype=bool)
-  """
   check_arraylike("jnp.linalg.qr", a)
   a, = promote_dtypes_inexact(jnp.asarray(a))
   if mode == "raw":
@@ -1262,44 +695,9 @@ def qr(a: ArrayLike, mode: str = "reduced") -> Array | QRResult:
   return QRResult(q, r)
 
 
+@implements(np.linalg.solve)
 @jit
 def solve(a: ArrayLike, b: ArrayLike) -> Array:
-  """Solve a linear system of equations
-
-  JAX implementation of :func:`numpy.linalg.solve`.
-
-  This solves a (batched) linear system of equations ``a @ x = b``
-  for ``x`` given ``a`` and ``b``.
-
-  Args:
-    a: array of shape ``(..., N, N)``.
-    b: array of shape ``(N,)`` (for 1-dimensional right-hand-side) or
-      ``(..., N, M)`` (for batched 2-dimensional right-hand-side).
-
-  Returns:
-    An array containing the result of the linear solve. The result has shape ``(..., N)``
-    if ``b`` is of shape ``(N,)``, and has shape ``(..., N, M)`` otherwise.
-
-  See also:
-    - :func:`jax.scipy.linalg.solve`: SciPy-style API for solving linear systems.
-    - :func:`jax.lax.custom_linear_solve`: matrix-free linear solver.
-
-  Example:
-    A simple 3x3 linear system:
-
-    >>> A = jnp.array([[1., 2., 3.],
-    ...                [2., 4., 2.],
-    ...                [3., 2., 1.]])
-    >>> b = jnp.array([14., 16., 10.])
-    >>> x = jnp.linalg.solve(A, b)
-    >>> x
-    Array([1., 2., 3.], dtype=float32)
-
-    Confirming that the result solves the system:
-
-    >>> jnp.allclose(A @ x, b)
-    Array(True, dtype=bool)
-  """
   check_arraylike("jnp.linalg.solve", a, b)
   a, b = promote_dtypes_inexact(jnp.asarray(a), jnp.asarray(b))
 
@@ -1364,82 +762,29 @@ def _lstsq(a: ArrayLike, b: ArrayLike, rcond: float | None, *,
 
 _jit_lstsq = jit(partial(_lstsq, numpy_resid=False))
 
+@implements(np.linalg.lstsq, lax_description=textwrap.dedent("""\
+    It has two important differences:
 
+    1. In `numpy.linalg.lstsq`, the default `rcond` is `-1`, and warns that in the future
+       the default will be `None`. Here, the default rcond is `None`.
+    2. In `np.linalg.lstsq` the returned residuals are empty for low-rank or over-determined
+       solutions. Here, the residuals are returned in all cases, to make the function
+       compatible with jit. The non-jit compatible numpy behavior can be recovered by
+       passing numpy_resid=True.
+
+    The lstsq function does not currently have a custom JVP rule, so the gradient is
+    poorly behaved for some inputs, particularly for low-rank `a`.
+    """))
 def lstsq(a: ArrayLike, b: ArrayLike, rcond: float | None = None, *,
           numpy_resid: bool = False) -> tuple[Array, Array, Array, Array]:
-  """
-  Return the least-squares solution to a linear equation.
-
-  JAX implementation of :func:`numpy.linalg.lstsq`.
-
-  Args:
-    a: array of shape ``(M, N)`` representing the coefficient matrix.
-    b: array of shape ``(M,)`` or ``(M, K)`` representing the right-hand side.
-    rcond: Cut-off ratio for small singular values. Singular values smaller than
-      ``rcond * largest_singular_value`` are treated as zero. If None (default),
-      the optimal value will be used to reduce floating point errors.
-    numpy_resid: If True, compute and return residuals in the same way as NumPy's
-      `linalg.lstsq`. This is necessary if you want to precisely replicate NumPy's
-      behavior. If False (default), a more efficient method is used to compute residuals.
-
-  Returns:
-    Tuple of arrays ``(x, resid, rank, s)`` where
-
-    - ``x`` is a shape ``(N,)`` or ``(N, K)`` array containing the least-squares solution.
-    - ``resid`` is the sum of squared residual of shape ``()`` or ``(K,)``.
-    - ``rank`` is the rank of the matrix ``a``.
-    - ``s`` is the singular values of the matrix ``a``.
-
-  Example:
-    >>> a = jnp.array([[1, 2],
-    ...                [3, 4]])
-    >>> b = jnp.array([5, 6])
-    >>> x, _, _, _ = jnp.linalg.lstsq(a, b)
-    >>> with jnp.printoptions(precision=3):
-    ...   print(x)
-    [-4.   4.5]
-  """
   check_arraylike("jnp.linalg.lstsq", a, b)
   if numpy_resid:
     return _lstsq(a, b, rcond, numpy_resid=True)
   return _jit_lstsq(a, b, rcond)
 
 
+@implements(getattr(np.linalg, "cross", None))
 def cross(x1: ArrayLike, x2: ArrayLike, /, *, axis=-1):
-  r"""Compute the corss-product of two 3D vectors
-
-  JAX implementation of :func:`numpy.linalg.cross`
-
-  Args:
-    x1: N-dimesional array, with ``x1.shape[axis] == 3``
-    x2: N-dimensional array, with ``x2.shape[axis] == 3``, and other axes
-      broadcast-compatible with ``x1``.
-    axis: axis along which to take the cross product (default: -1).
-
-  Returns:
-    array containing the result of the cross-product
-
-  See Also:
-    :func:`jax.numpy.cross`: more flexible cross-product API.
-
-  Example:
-
-    Showing that :math:`\hat{x} \times \hat{y} = \hat{z}`:
-
-    >>> x = jnp.array([1., 0., 0.])
-    >>> y = jnp.array([0., 1., 0.])
-    >>> jnp.linalg.cross(x, y)
-    Array([0., 0., 1.], dtype=float32)
-
-    Cross product of :math:`\hat{x}` with all three standard unit vectors,
-    via broadcasting:
-
-    >>> xyz = jnp.eye(3)
-    >>> jnp.linalg.cross(x, xyz, axis=-1)
-    Array([[ 0.,  0.,  0.],
-           [ 0.,  0.,  1.],
-           [ 0., -1.,  0.]], dtype=float32)
-  """
   check_arraylike("jnp.linalg.outer", x1, x2)
   x1, x2 = jnp.asarray(x1), jnp.asarray(x2)
   if x1.shape[axis] != 3 or x2.shape[axis] != 3:
@@ -1450,29 +795,8 @@ def cross(x1: ArrayLike, x2: ArrayLike, /, *, axis=-1):
   return jnp.cross(x1, x2, axis=axis)
 
 
+@implements(getattr(np.linalg, "outer", None))
 def outer(x1: ArrayLike, x2: ArrayLike, /) -> Array:
-  """Compute the outer product of two 1-dimensional arrays.
-
-  JAX implementation of :func:`numpy.linalg.outer`.
-
-  Args:
-    x1: array
-    x2: array
-
-  Returns:
-    array containing the outer product of ``x1`` and ``x2``
-
-  See also:
-    :func:`jax.numpy.outer`: similar function in the main :mod:`jax.numpy` module.
-
-  Example:
-    >>> x1 = jnp.array([1, 2, 3])
-    >>> x2 = jnp.array([4, 5, 6])
-    >>> jnp.linalg.outer(x1, x2)
-    Array([[ 4,  5,  6],
-           [ 8, 10, 12],
-           [12, 15, 18]], dtype=int32)
-  """
   check_arraylike("jnp.linalg.outer", x1, x2)
   x1, x2 = jnp.asarray(x1), jnp.asarray(x2)
   if x1.ndim != 1 or x2.ndim != 1:
@@ -1480,83 +804,18 @@ def outer(x1: ArrayLike, x2: ArrayLike, /) -> Array:
   return x1[:, None] * x2[None, :]
 
 
+@implements(getattr(np.linalg, "matrix_norm", None))
 def matrix_norm(x: ArrayLike, /, *, keepdims: bool = False, ord: str = 'fro') -> Array:
-  """Compute the norm of a matrix or stack of matrices.
-
-  JAX implementation of :func:`numpy.linalg.matrix_norm`
-
-  Args:
-    x: array of shape ``(..., M, N)`` for which to take the norm.
-    keepdims: if True, keep the reduced dimensions in the output.
-    ord: A string or int specifying the type of norm; default is the Frobenius norm.
-      See :func:`numpy.linalg.norm` for details on available options.
-
-  Returns:
-    array containing the norm of ``x``. Has shape ``x.shape[:-2]`` if ``keepdims`` is
-    False, or shape ``(..., 1, 1)`` if ``keepdims`` is True.
-
-  See also:
-    - :func:`jax.numpy.linalg.vector_norm`: Norm of a vector or stack of vectors.
-    - :func:`jax.numpy.linalg.norm`: More general matrix or vector norm.
-
-  Examples:
-    >>> x = jnp.array([[1, 2, 3],
-    ...                [4, 5, 6],
-    ...                [7, 8, 9]])
-    >>> jnp.linalg.matrix_norm(x)
-    Array(16.881943, dtype=float32)
+  """
+  Computes the matrix norm of a matrix (or a stack of matrices) x.
   """
   check_arraylike('jnp.linalg.matrix_norm', x)
   return norm(x, ord=ord, keepdims=keepdims, axis=(-2, -1))
 
 
+@implements(getattr(np.linalg, "matrix_transpose", None))
 def matrix_transpose(x: ArrayLike, /) -> Array:
-  """Transpose a matrix or stack of matrices.
-
-  JAX implementation of :func:`numpy.linalg.matrix_transpose`.
-
-  Args:
-    x: array of shape ``(..., M, N)``
-
-  Returns:
-    array of shape ``(..., N, M)`` containing the matrix transpose of ``x``.
-
-  See also:
-    :func:`jax.numpy.transpose`: more general transpose operation.
-
-  Examples:
-    Transpose of a single matrix:
-
-    >>> x = jnp.array([[1, 2, 3],
-    ...                [4, 5, 6]])
-    >>> jnp.linalg.matrix_transpose(x)
-    Array([[1, 4],
-           [2, 5],
-           [3, 6]], dtype=int32)
-
-    Transpose of a stack of matrices:
-
-    >>> x = jnp.array([[[1, 2],
-    ...                 [3, 4]],
-    ...                [[5, 6],
-    ...                 [7, 8]]])
-    >>> jnp.linalg.matrix_transpose(x)
-    Array([[[1, 3],
-            [2, 4]],
-    <BLANKLINE>
-           [[5, 7],
-            [6, 8]]], dtype=int32)
-
-    For convenience, the same computation can be done via the
-    :attr:`~jax.Array.mT` property of JAX array objects:
-
-    >>> x.mT
-    Array([[[1, 3],
-            [2, 4]],
-    <BLANKLINE>
-           [[5, 7],
-            [6, 8]]], dtype=int32)
-  """
+  """Transposes a matrix (or a stack of matrices) x."""
   check_arraylike('jnp.linalg.matrix_transpose', x)
   x_arr = jnp.asarray(x)
   ndim = x_arr.ndim
@@ -1565,41 +824,10 @@ def matrix_transpose(x: ArrayLike, /) -> Array:
   return jax.lax.transpose(x_arr, (*range(ndim - 2), ndim - 1, ndim - 2))
 
 
+@implements(getattr(np.linalg, "vector_norm", None))
 def vector_norm(x: ArrayLike, /, *, axis: int | None = None, keepdims: bool = False,
                 ord: int | str = 2) -> Array:
-  """Computes the vector norm of a vector or batch of vectors.
-
-  JAX implementation of :func:`numpy.linalg.vector_norm`.
-
-  Args:
-    x: N-dimensional array for which to take the norm.
-    axis: optional axis along which to compute the vector norm. If None (default)
-      then ``x`` is flattened and the norm is taken over all values.
-    keepdims: if True, keep the reduced dimensions in the output.
-    ord: A string or int specifying the type of norm; default is the 2-norm.
-      See :func:`numpy.linalg.norm` for details on available options.
-
-  Returns:
-    array containing the norm of ``x``.
-
-  See also:
-    - :func:`jax.numpy.linalg.matrix_norm`: Norm of a matrix or stack of matrices.
-    - :func:`jax.numpy.linalg.norm`: More general matrix or vector norm.
-
-  Examples:
-    Norm of a single vector:
-
-    >>> x = jnp.array([1., 2., 3.])
-    >>> jnp.linalg.vector_norm(x)
-    Array(3.7416575, dtype=float32)
-
-    Norm of a batch of vectors:
-
-    >>> x = jnp.array([[1., 2., 3.],
-    ...                [4., 5., 7.]])
-    >>> jnp.linalg.vector_norm(x, axis=1)
-    Array([3.7416575, 9.486833 ], dtype=float32)
-  """
+  """Computes the vector norm of a vector (or batch of vectors) x."""
   check_arraylike('jnp.linalg.vector_norm', x)
   if axis is None:
     result = norm(jnp.ravel(x), ord=ord)
@@ -1609,226 +837,31 @@ def vector_norm(x: ArrayLike, /, *, axis: int | None = None, keepdims: bool = Fa
   return norm(x, axis=axis, keepdims=keepdims, ord=ord)
 
 
+@implements(getattr(np.linalg, "vecdot", None))
 def vecdot(x1: ArrayLike, x2: ArrayLike, /, *, axis: int = -1) -> Array:
-  """Compute the (batched) vector dot product of two arrays.
-
-  JAX implementation of :func:`numpy.linalg.vecdot`.
-
-  Args:
-    x1: left-hand side array.
-    x2: right-hand side array. Size of ``x2[axis]`` must match size of ``x1[axis]``.
-    axis: axis along which to compute the dot product (default: -1)
-
-  Returns:
-    array containing the dot product of ``x1`` and ``x2`` along ``axis``. The
-    non-contracted dimensions are broadcast together.
-
-  See also:
-    - :func:`jax.numpy.vecdot`: similar API in the ``jax.numpy`` namespace.
-    - :func:`jax.numpy.linalg.matmul`: matrix multiplication.
-    - :func:`jax.numpy.linalg.tensordot`: general tensor dot product.
-
-  Examples:
-    Vector dot product of two 1D arrays:
-
-    >>> x1 = jnp.array([1, 2, 3])
-    >>> x2 = jnp.array([4, 5, 6])
-    >>> jnp.linalg.vecdot(x1, x2)
-    Array(32, dtype=int32)
-
-
-    Batched vector dot product of two 2D arrays:
-
-    >>> x1 = jnp.array([[1, 2, 3],
-    ...                 [4, 5, 6]])
-    >>> x2 = jnp.array([[2, 3, 4]])
-    >>> jnp.linalg.vecdot(x1, x2, axis=-1)
-    Array([20, 47], dtype=int32)
-  """
-  check_arraylike('jnp.linalg.vecdot', x1, x2)
   return jnp.vecdot(x1, x2, axis=axis)
 
 
+@implements(getattr(np.linalg, "matmul", None))
 def matmul(x1: ArrayLike, x2: ArrayLike, /) -> Array:
-  """Perform a matrix multiplication.
-
-  JAX implementation of :func:`numpy.linalg.matmul`.
-
-  Args:
-    x1: first input array, of shape ``(..., N)``.
-    x2: second input array. Must have shape ``(N,)`` or ``(..., M, N)``.
-      In the multi-dimensional case, leading dimensions must be broadcast-compatible
-      with the leading dimensions of ``x1``.
-
-  Returns:
-    array containing the matrix product of the inputs. Shape is ``x1.shape[:-1]``
-    if ``x2.ndim == 1``, otherwise the shape is ``(..., M)``.
-
-  See Also:
-    :func:`jax.numpy.matmul`: NumPy API for this function.
-    :func:`jax.numpy.linalg.vecdot`: batched vector product.
-    :func:`jax.numpy.linalg.tensordot`: batched tensor product.
-
-  Examples:
-    Vector dot products:
-
-    >>> x1 = jnp.array([1, 2, 3])
-    >>> x2 = jnp.array([4, 5, 6])
-    >>> jnp.linalg.matmul(x1, x2)
-    Array(32, dtype=int32)
-
-    Matrix dot product:
-
-    >>> x1 = jnp.array([[1, 2, 3],
-    ...                 [4, 5, 6]])
-    >>> x2 = jnp.array([[1, 2],
-    ...                 [3, 4],
-    ...                 [5, 6]])
-    >>> jnp.linalg.matmul(x1, x2)
-    Array([[22, 28],
-           [49, 64]], dtype=int32)
-
-    For convenience, in all cases you can do the same computation using
-    the ``@`` operator:
-
-    >>> x1 @ x2
-    Array([[22, 28],
-           [49, 64]], dtype=int32)
-  """
   check_arraylike('jnp.linalg.matmul', x1, x2)
   return jnp.matmul(x1, x2)
 
 
+@implements(getattr(np.linalg, "tensordot", None))
 def tensordot(x1: ArrayLike, x2: ArrayLike, /, *,
               axes: int | tuple[Sequence[int], Sequence[int]] = 2) -> Array:
-  """Compute the tensor dot product of two N-dimensional arrays.
-
-  JAX implementation of :func:`numpy.linalg.tensordot`.
-
-  Args:
-    x1: N-dimensional array
-    x2: M-dimensional array
-    axes: integer or tuple of sequences of integers. If an integer `k`, then
-      sum over the last `k` axes of ``x1`` and the first `k` axes of ``x2``,
-      in order. If a tuple, then ``axes[0]`` specifies the axes of ``x1`` and
-      ``axes[1]`` specifies the axes of ``x2``.
-
-  Returns:
-    array containing the tensor dot product of the inputs
-
-  See also:
-    :func:`jax.numpy.tensordot`: equivalent API in the :mod:`jax.numpy` namespace.
-    :func:`jax.numpy.einsum`: NumPy API for more general tensor contractions.
-    :func:`jax.lax.dot_general`: XLA API for more general tensor contractions.
-
-  Examples:
-    >>> x1 = jnp.arange(24.).reshape(2, 3, 4)
-    >>> x2 = jnp.ones((3, 4, 5))
-    >>> jnp.tensordot(x1, x2)
-    Array([[ 66.,  66.,  66.,  66.,  66.],
-           [210., 210., 210., 210., 210.]], dtype=float32)
-
-    Equivalent result when specifying the axes as explicit sequences:
-
-    >>> jnp.linalg.tensordot(x1, x2, axes=([1, 2], [0, 1]))
-    Array([[ 66.,  66.,  66.,  66.,  66.],
-           [210., 210., 210., 210., 210.]], dtype=float32)
-
-    Equivalent result via :func:`~jax.numpy.einsum`:
-
-    >>> jnp.einsum('ijk,jkm->im', x1, x2)
-    Array([[ 66.,  66.,  66.,  66.,  66.],
-           [210., 210., 210., 210., 210.]], dtype=float32)
-
-    Setting ``axes=1`` for two-dimensional inputs is equivalent to a matrix
-    multiplication:
-
-    >>> x1 = jnp.array([[1, 2],
-    ...                 [3, 4]])
-    >>> x2 = jnp.array([[1, 2, 3],
-    ...                 [4, 5, 6]])
-    >>> jnp.linalg.tensordot(x1, x2, axes=1)
-    Array([[ 9, 12, 15],
-           [19, 26, 33]], dtype=int32)
-    >>> x1 @ x2
-    Array([[ 9, 12, 15],
-           [19, 26, 33]], dtype=int32)
-
-    Setting ``axes=0`` for one-dimensional inputs is equivalent to
-    ``jnp.linalg.outer``:
-
-    >>> x1 = jnp.array([1, 2])
-    >>> x2 = jnp.array([1, 2, 3])
-    >>> jnp.linalg.tensordot(x1, x2, axes=0)
-    Array([[1, 2, 3],
-           [2, 4, 6]], dtype=int32)
-    >>> jnp.linalg.outer(x1, x2)
-    Array([[1, 2, 3],
-           [2, 4, 6]], dtype=int32)
-  """
   check_arraylike('jnp.linalg.tensordot', x1, x2)
   return jnp.tensordot(x1, x2, axes=axes)
 
 
+@implements(getattr(np.linalg, "svdvals", None))
 def svdvals(x: ArrayLike, /) -> Array:
-  """Compute the singular values of a matrix.
-
-  JAX implementation of :func:`numpy.linalg.svdvals`.
-
-  Args:
-    x: array of shape ``(..., M, N)`` for which singular values will be computed.
-
-  Returns:
-    array of singular values of shape ``(..., K)`` with ``K = min(M, N)``.
-
-  See also:
-    :func:`jax.numpy.linalg.svd`: compute singular values and singular vectors
-
-  Example:
-    >>> x = jnp.array([[1, 2, 3],
-    ...                [4, 5, 6]])
-    >>> jnp.linalg.svdvals(x)
-    Array([9.508031 , 0.7728694], dtype=float32)
-  """
   check_arraylike('jnp.linalg.svdvals', x)
   return svd(x, compute_uv=False, hermitian=False)
 
 
+@implements(getattr(np.linalg, "diagonal", None))
 def diagonal(x: ArrayLike, /, *, offset: int = 0) -> Array:
-  """Extract the diagonal of an matrix or stack of matrices.
-
-  JAX implementation of :func:`numpy.linalg.diagonal`.
-
-  Args:
-    x: array of shape ``(..., M, N)`` from which the diagonal will be extracted.
-    offset: positive or negative offset from the main diagonal.
-
-  Returns:
-    Array of shape ``(..., K)`` where ``K`` is the length of the specified diagonal.
-
-  See Also:
-    - :func:`jax.numpy.diagonal`: more general functionality for extracting diagonals.
-    - :func:`jax.numpy.diag`: create a diagonal matrix from values.
-
-  Examples:
-    Diagonals of a single matrix:
-
-    >>> x = jnp.array([[1,  2,  3,  4],
-    ...                [5,  6,  7,  8],
-    ...                [9, 10, 11, 12]])
-    >>> jnp.linalg.diagonal(x)
-    Array([ 1,  6, 11], dtype=int32)
-    >>> jnp.linalg.diagonal(x, offset=1)
-    Array([ 2,  7, 12], dtype=int32)
-    >>> jnp.linalg.diagonal(x, offset=-1)
-    Array([ 5, 10], dtype=int32)
-
-    Batched diagonals:
-
-    >>> x = jnp.arange(24).reshape(2, 3, 4)
-    >>> jnp.linalg.diagonal(x)
-    Array([[ 0,  5, 10],
-           [12, 17, 22]], dtype=int32)
-  """
   check_arraylike('jnp.linalg.diagonal', x)
   return jnp.diagonal(x, offset=offset, axis1=-2, axis2=-1)
diff --git a/jax/_src/pretty_printer.py b/jax/_src/pretty_printer.py
index ec5c34ab8..0614bb8a8 100644
--- a/jax/_src/pretty_printer.py
+++ b/jax/_src/pretty_printer.py
@@ -31,7 +31,7 @@ from collections.abc import Sequence
 import enum
 from functools import partial
 import sys
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 from jax._src import config
 from jax._src import util
@@ -69,12 +69,23 @@ CAN_USE_COLOR = _can_use_color()
 class Doc(util.StrictABC):
   __slots__ = ()
 
-  def format(self, width: int = 80, use_color: bool | None = None,
-             annotation_prefix=" # ") -> str:
+  def format(
+    self, width: int = 80, *, use_color: bool | None = None,
+    annotation_prefix: str = " # ",
+    source_map: list[list[tuple[int, int, Any]]] | None = None
+  ) -> str:
+    """
+    Formats a pretty-printer document as a string.
+
+    Args:
+    source_map: for each line in the output, contains a list of
+      (start column, end column, source) tuples. Each tuple associates a
+      region of output text with a source.
+    """
     if use_color is None:
       use_color = CAN_USE_COLOR and _PPRINT_USE_COLOR.value
     return _format(self, width, use_color=use_color,
-                   annotation_prefix=annotation_prefix)
+                   annotation_prefix=annotation_prefix, source_map=source_map)
 
   def __str__(self):
     return self.format()
@@ -147,6 +158,21 @@ class _NestDoc(Doc):
   def __repr__(self): return f"nest({self.n, self.child})"
 
 
+_NO_SOURCE = object()
+
+class _SourceMapDoc(Doc):
+  __slots__ = ("child", "source")
+  child: Doc
+  source: Any
+
+  def __init__(self, child: Doc, source: Any):
+    assert isinstance(child, Doc), child
+    self.child = child
+    self.source = source
+
+  def __repr__(self): return f"source({self.child}, {self.source})"
+
+
 Color = enum.Enum("_Color", ["BLACK", "RED", "GREEN", "YELLOW", "BLUE",
                              "MAGENTA", "CYAN", "WHITE", "RESET"])
 Intensity = enum.Enum("_Intensity", ["DIM", "NORMAL", "BRIGHT"])
@@ -193,7 +219,7 @@ def _fits(doc: Doc, width: int, agenda: list[tuple[int, _BreakMode, Doc]]
       agenda.append((i + doc.n, m, doc.child))
     elif isinstance(doc, _GroupDoc):
       agenda.append((i, _BreakMode.FLAT, doc.child))
-    elif isinstance(doc, _ColorDoc):
+    elif isinstance(doc, _ColorDoc) or isinstance(doc, _SourceMapDoc):
       agenda.append((i, m, doc.child))
     else:
       raise ValueError("Invalid document ", doc)
@@ -224,7 +250,7 @@ def _sparse(doc: Doc) -> bool:
       agenda.append(doc.child)
     elif isinstance(doc, _GroupDoc):
       agenda.append(doc.child)
-    elif isinstance(doc, _ColorDoc):
+    elif isinstance(doc, _ColorDoc) or isinstance(doc, _SourceMapDoc):
       agenda.append(doc.child)
     else:
       raise ValueError("Invalid document ", doc)
@@ -241,6 +267,7 @@ class _State(NamedTuple):
   mode: _BreakMode
   doc: Doc
   color: _ColorState
+  source_map: Any
 
 class _Line(NamedTuple):
   text: str
@@ -283,17 +310,29 @@ def _align_annotations(lines):
 
 
 
-def _format(doc: Doc, width: int, *, use_color, annotation_prefix) -> str:
+def _format(
+  doc: Doc, width: int, *, use_color: bool, annotation_prefix: str,
+  source_map: list[list[tuple[int, int, Any]]] | None
+) -> str:
   lines = []
   default_colors = _ColorState(Color.RESET, Color.RESET, Intensity.NORMAL)
   annotation_colors = _ColorState(Color.RESET, Color.RESET, Intensity.DIM)
   color_state = default_colors
-  agenda = [_State(0, _BreakMode.BREAK, doc, default_colors)]
+  source_start = 0       # The column at which the current source region starts.
+  source = _NO_SOURCE    # The currently active source region.
+  line_source_map = []  # Source maps for the current line of text.
+  agenda = [_State(0, _BreakMode.BREAK, doc, default_colors, source)]
   k = 0
   line_text = ""
   line_annotations = []
   while len(agenda) > 0:
-    i, m, doc, color = agenda.pop()
+    i, m, doc, color, agenda_source = agenda.pop()
+    if source_map is not None and agenda_source != source:
+      pos = len(line_text)
+      if source_start != pos and source is not _NO_SOURCE:
+        line_source_map.append((source_start, pos, source))
+      source = agenda_source
+      source_start = pos
     if isinstance(doc, _NilDoc):
       pass
     elif isinstance(doc, _TextDoc):
@@ -304,7 +343,7 @@ def _format(doc: Doc, width: int, *, use_color, annotation_prefix) -> str:
         line_annotations.append(doc.annotation)
       k += len(doc.text)
     elif isinstance(doc, _ConcatDoc):
-      agenda.extend(_State(i, m, d, color)
+      agenda.extend(_State(i, m, d, color, source)
                     for d in reversed(doc.children))
     elif isinstance(doc, _BreakDoc):
       if m == _BreakMode.BREAK:
@@ -313,6 +352,13 @@ def _format(doc: Doc, width: int, *, use_color, annotation_prefix) -> str:
                                                  annotation_colors)
           line_text += color_str
         lines.append(_Line(line_text, k, line_annotations))
+        if source_map is not None:
+          pos = len(line_text)
+          if source_start != pos and source is not _NO_SOURCE:
+            line_source_map.append((source_start, pos, source))
+          source_map.append(line_source_map)
+          line_source_map = []
+          source_start = i
         line_text = " " * i
         line_annotations = []
         k = i
@@ -322,20 +368,22 @@ def _format(doc: Doc, width: int, *, use_color, annotation_prefix) -> str:
         line_text += doc.text
         k += len(doc.text)
     elif isinstance(doc, _NestDoc):
-      agenda.append(_State(i + doc.n, m, doc.child, color))
+      agenda.append(_State(i + doc.n, m, doc.child, color, source))
     elif isinstance(doc, _GroupDoc):
       # In Lindig's paper, _fits is passed the remainder of the document.
       # I'm pretty sure that's a bug and we care only if the current group fits!
       if (_sparse(doc)
           and _fits(doc, width - k, [(i, _BreakMode.FLAT, doc.child)])):
-        agenda.append(_State(i, _BreakMode.FLAT, doc.child, color))
+        agenda.append(_State(i, _BreakMode.FLAT, doc.child, color, source))
       else:
-        agenda.append(_State(i, _BreakMode.BREAK, doc.child, color))
+        agenda.append(_State(i, _BreakMode.BREAK, doc.child, color, source))
     elif isinstance(doc, _ColorDoc):
       color = _ColorState(doc.foreground or color.foreground,
                           doc.background or color.background,
                           doc.intensity or color.intensity)
-      agenda.append(_State(i, m, doc.child, color))
+      agenda.append(_State(i, m, doc.child, color, source))
+    elif isinstance(doc, _SourceMapDoc):
+      agenda.append(_State(i, m, doc.child, color, doc.source))
     else:
       raise ValueError("Invalid document ", doc)
 
@@ -343,6 +391,11 @@ def _format(doc: Doc, width: int, *, use_color, annotation_prefix) -> str:
     color_state, color_str = _update_color(use_color, color_state,
                                            annotation_colors)
     line_text += color_str
+  if source_map is not None:
+    pos = len(line_text)
+    if source_start != pos and source is not _NO_SOURCE:
+      line_source_map.append((source_start, pos, source))
+    source_map.append(line_source_map)
   lines.append(_Line(line_text, k, line_annotations))
   lines = _align_annotations(lines)
   out = "\n".join(
@@ -406,6 +459,17 @@ def color(doc: Doc, *, foreground: Color | None = None,
                    intensity=intensity)
 
 
+def source_map(doc: Doc, source: Any):
+  """Source mapping.
+
+  A source map associates a region of the pretty-printer's text output with a
+  source location that produced it. For the purposes of the pretty printer a
+  ``source`` may be any object: we require only that we can compare sources for
+  equality. A text region to source object mapping can be populated as a side
+  output of the ``format`` method.
+  """
+  return _SourceMapDoc(doc, source)
+
 type_annotation = partial(color, intensity=Intensity.NORMAL,
                           foreground=Color.MAGENTA)
 keyword = partial(color, intensity=Intensity.BRIGHT, foreground=Color.BLUE)
diff --git a/jax/tools/toolchains/BUILD b/jax/tools/toolchains/BUILD
deleted file mode 100644
index 1401ee12a..000000000
--- a/jax/tools/toolchains/BUILD
+++ /dev/null
@@ -1,30 +0,0 @@
-# Copyright 2024 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-licenses(["notice"])
-
-package(
-    default_applicable_licenses = [],
-    default_visibility = ["//visibility:public"],
-)
-
-# Refer to https://bazel.build/configure/windows#clang for potential changes.
-platform(
-    name = "x64_windows-clang-cl",
-    constraint_values = [
-        "@platforms//cpu:x86_64",
-        "@platforms//os:windows",
-        "@bazel_tools//tools/cpp:clang-cl",
-    ],
-)


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index f0b8f7570..0dfeda8a2 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1477,6 +1477,15 @@ jax_test(
     ],
 )
 
+py_test(
+    name = "pretty_printer_test",
+    srcs = ["pretty_printer_test.py"],
+    deps = [
+        "//jax",
+        "//jax:test_util",
+    ],
+)
+
 exports_files(
     [
         "api_test.py",
diff --git a/tests/pretty_printer_test.py b/tests/pretty_printer_test.py
new file mode 100644
index 000000000..d87708c9d
--- /dev/null
+++ b/tests/pretty_printer_test.py
@@ -0,0 +1,36 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import absltest
+
+from jax._src import test_util as jtu
+from jax._src import pretty_printer as pp
+
+
+class PrettyPrinterTest(jtu.JaxTestCase):
+
+  def testSourceMap(self):
+    doc = pp.concat([
+      pp.text("abc"), pp.source_map(pp.text("def"), 101),
+      pp.source_map(pp.concat([pp.text("gh"), pp.brk(""), pp.text("ijkl")]), 77),
+      pp.text("mn"),
+    ])
+    source_map = []
+    out = doc.format(width=8, source_map=source_map)
+    self.assertEqual(out, "abcdefgh\nijklmn")
+    self.assertEqual(source_map, [[(3, 6, 101), (6, 8, 77)], [(0, 4, 77)]])
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
index 39c4a0ef9..fa21894de 100644
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -65,9 +65,6 @@ class LaxRandomTest(jtu.JaxTestCase):
     # whether RBG keys may be involved, but that's no longer exact.
     if config.enable_custom_prng.value and samples.dtype == jnp.bfloat16:
       return
-    # kstest does not understand bfloat16 input, so cast to float32.
-    if samples.dtype == jnp.bfloat16:
-      samples = samples.astype('float32')
     # kstest fails for infinities starting in scipy 1.12
     # (https://github.com/scipy/scipy/issues/20386)
     # TODO(jakevdp): remove this logic if/when fixed upstream.