Reorganize linear algebra (#2037)

* Move matrix decomposition Ops into a sub-module

* Move solve Ops into a sub-module

* Move misc solve Ops

* Move inverse Ops to sub-module

* Move matrix product Ops to sub-module

* Move matrix construction Ops to sub-module

* Move matrix summarization Ops to sub-module

* Re-import all public functions via `linalg.py`

* Issue depreciation warning when importing from nlinalg/slinalg

* Update rewrite imports to use _linalg locations

* Update link dispatch imports to use _linalg locations

* Update numba dispatch imports to use _linalg locations

* Organize linalg rewrites according to library

* Organize linalg dispatches according to library

* Organize numba dispatches according to library

* Organize linalg tests according to library

* Organize linalg rewrite tests according to library

* Rename tests/tensor/test_linalg -> tests/tensor/linalg

* Reorganize dispatch test according to library

* Clean up stale linalg imports

* Update docs

* Rename tests/link/*/test_linalg -> tests/link/*/linalg

* Remove stray empty test_linalg/__init__.py under numba linalg

* Update numba link CI part to new linalg/ directory

* Move solve rewrites under tensor/rewriting/linalg

The scan-tracking solve rewrites previously lived in
pytensor/tensor/_linalg/solve/rewriting.py and were registered via a
side-effect import chain from tensor/_linalg/__init__.py. That chain
was broken in the linalg reorg (the _linalg package init was emptied),
and registration only kept working incidentally because
tensor/linalg.py happens to import from _linalg.solve.

Move them into tensor/rewriting/linalg/solve.py alongside the other
solve rewrites. The imports also needed updating: the v3 file imported
via tensor.slinalg, which is now the deprecation shim and would emit
DeprecationWarnings on every load.

Also drop the speculative __all__ and inv_to_solve re-export from
tensor/rewriting/linalg/__init__.py; the one test that imported
inv_to_solve from there now imports it from .inverse directly.

* Stop importing from deprecated nlinalg/slinalg in unrelated tests

These three test files happened to import individual Ops/functions from
pytensor.tensor.{nlinalg,slinalg}. Since those modules are now
deprecation shims, every test run was emitting DeprecationWarnings for
imports that have nothing to do with testing the shim itself.

Retarget the imports to the new pytensor.tensor._linalg.* locations,
matching the convention used by the new tests/tensor/linalg/ tests.

* Stop importing from deprecated nlinalg/slinalg in legacy test suites

tests/tensor/test_nlinalg.py and tests/tensor/test_slinalg.py were left
untouched by the linalg reorg and still imported individual symbols
from pytensor.tensor.{nlinalg,slinalg}. With those modules now being
deprecation shims, every collection of these files emitted dozens of
DeprecationWarnings.

These files are not testing the deprecation shim itself — they're the
original Op test suites. Retarget the imports to the new
pytensor.tensor._linalg.* locations.

* Add targeted test for nlinalg/slinalg deprecation shim

Asserts that accessing a moved name on the deprecated module emits a
DeprecationWarning, that the shim forwards to the same object exposed
by the new pytensor.tensor.linalg public API, and that unknown names
still raise AttributeError. Also fixes a stale doc comment in
_linalg/summary.py that pointed at the old nlinalg path.

* Drop automodule from deprecated nlinalg/slinalg doc pages

The deprecation pages still carried `.. automodule:: pytensor.tensor.{nlinalg,slinalg} :members:` directives. Each automodule
walks the module's `__dir__`, which on the new shims returns every
moved name and triggers `__getattr__` per name — emitting a
DeprecationWarning per attribute on every doc build, and producing
duplicate doc entries that already exist on the new
:ref:`libdoc_linalg` page.

Drop the automodule blocks; the pages are now pure deprecation notices
that point at `pytensor.tensor.linalg`.

* Group linalg API docs by category and pin them to __all__

The linalg.rst page used `.. automodule:: pytensor.tensor.linalg
:members:`, which renders all 36 functions from `__all__` as one
unstructured alphabetical list. Replace it with hand-grouped sections
(Constructors / Decomposition / Inverse / Products / Solve / Summary)
that mirror the internal `_linalg/` package layout, using explicit
`autofunction::` directives.

To stop the rst from drifting away from the public API, add
`tests/tensor/linalg/test_doc_api.py`. It parses every
`.. autofunction:: pytensor.tensor.linalg.<name>` directive out of the
rst and asserts the resulting set is exactly equal to
`pytensor.tensor.linalg.__all__`. Any future addition or removal on
either side fails the test with a precise diff.

Op classes remain undocumented in the rst — the public surface is
intentionally functions only.

* Tweak linalg directory structure

---------

Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

Author: jessegrabowski

.github/workflows/test.yml

@@ -101,11 +101,11 @@ jobs:
             default-mode: "CVM"
             python-version: "3.12"
             os: "ubuntu-latest"
-          - part: ["numba link", "tests/link/numba --ignore=tests/link/numba/test_slinalg.py"]
+          - part: ["numba link", "tests/link/numba --ignore=tests/link/numba/linalg"]
             default-mode: "CVM"
             python-version: "3.12"
             os: "ubuntu-latest"
-          - part: ["numba link slinalg", "tests/link/numba/test_slinalg.py"]
+          - part: ["numba link linalg", "tests/link/numba/linalg"]
             default-mode: "CVM"
             python-version: "3.13"
             os: "ubuntu-latest"

doc/library/tensor/index.rst

@@ -24,6 +24,7 @@ symbolic expressions using calls that look just like numpy calls, such as
     elemwise
     extra_ops
     io
+    linalg
     slinalg
     nlinalg
     fft

doc/library/tensor/linalg.rst

@@ -0,0 +1,72 @@
+.. _libdoc_linalg:
+
+===================================================================
+:mod:`tensor.linalg` --  Linear Algebra Operations
+===================================================================
+
+.. module:: tensor.linalg
+   :platform: Unix, Windows
+   :synopsis: Linear Algebra Operations
+
+The :mod:`pytensor.tensor.linalg` module exposes the user-facing linear
+algebra API.
+
+Constructors
+============
+
+.. autofunction:: pytensor.tensor.linalg.block_diag
+
+Decomposition
+=============
+
+.. autofunction:: pytensor.tensor.linalg.cholesky
+.. autofunction:: pytensor.tensor.linalg.lu
+.. autofunction:: pytensor.tensor.linalg.lu_factor
+.. autofunction:: pytensor.tensor.linalg.pivot_to_permutation
+.. autofunction:: pytensor.tensor.linalg.qr
+.. autofunction:: pytensor.tensor.linalg.svd
+.. autofunction:: pytensor.tensor.linalg.eig
+.. autofunction:: pytensor.tensor.linalg.eigh
+.. autofunction:: pytensor.tensor.linalg.eigvalsh
+.. autofunction:: pytensor.tensor.linalg.schur
+.. autofunction:: pytensor.tensor.linalg.qz
+.. autofunction:: pytensor.tensor.linalg.ordqz
+
+Inverse
+=======
+
+.. autofunction:: pytensor.tensor.linalg.inv
+.. autofunction:: pytensor.tensor.linalg.pinv
+.. autofunction:: pytensor.tensor.linalg.tensorinv
+
+Products
+========
+
+.. autofunction:: pytensor.tensor.linalg.kron
+.. autofunction:: pytensor.tensor.linalg.matrix_dot
+.. autofunction:: pytensor.tensor.linalg.matrix_power
+.. autofunction:: pytensor.tensor.linalg.expm
+
+Solve
+=====
+
+.. autofunction:: pytensor.tensor.linalg.solve
+.. autofunction:: pytensor.tensor.linalg.solve_triangular
+.. autofunction:: pytensor.tensor.linalg.lu_solve
+.. autofunction:: pytensor.tensor.linalg.cho_solve
+.. autofunction:: pytensor.tensor.linalg.lstsq
+.. autofunction:: pytensor.tensor.linalg.tensorsolve
+.. autofunction:: pytensor.tensor.linalg.tridiagonal_lu_factor
+.. autofunction:: pytensor.tensor.linalg.tridiagonal_lu_solve
+.. autofunction:: pytensor.tensor.linalg.solve_continuous_lyapunov
+.. autofunction:: pytensor.tensor.linalg.solve_discrete_lyapunov
+.. autofunction:: pytensor.tensor.linalg.solve_discrete_are
+.. autofunction:: pytensor.tensor.linalg.solve_sylvester
+
+Summary
+=======
+
+.. autofunction:: pytensor.tensor.linalg.det
+.. autofunction:: pytensor.tensor.linalg.slogdet
+.. autofunction:: pytensor.tensor.linalg.norm
+.. autofunction:: pytensor.tensor.linalg.trace

doc/library/tensor/nlinalg.rst

@@ -1,22 +1,16 @@
-..  ../../../../pytensor/sandbox/nlinalg.py
-
-.. _libdoc_linalg:
+.. _libdoc_nlinalg:
 
 ===================================================================
-:mod:`tensor.nlinalg` --  Linear Algebra Ops Using Numpy
+:mod:`tensor.nlinalg` --  Linear Algebra Ops Using Numpy (deprecated)
 ===================================================================
 
 .. module:: tensor.nlinalg
    :platform: Unix, Windows
-   :synopsis: Linear Algebra Ops Using Numpy
+   :synopsis: Linear Algebra Ops Using Numpy (deprecated)
 .. moduleauthor:: LISA
 
-.. note::
-
-   This module is not imported by default. You need to import it to use it.
-
-API
-===
-
-.. automodule:: pytensor.tensor.nlinalg
-    :members:
+.. deprecated:: 2.x
+   The ``nlinalg`` module is deprecated. Use :mod:`pytensor.tensor.linalg` instead.
+   All symbols previously exported from ``nlinalg`` are available from
+   ``pytensor.tensor.linalg`` — see :ref:`libdoc_linalg`. Imports from
+   ``nlinalg`` will be removed in PyTensor 3.0.

doc/library/tensor/slinalg.rst

@@ -1,23 +1,16 @@
-..  ../../../../pytensor/sandbox/slinalg.py
-
 .. _libdoc_slinalg:
 
 ===================================================================
-:mod:`tensor.slinalg` --  Linear Algebra Ops Using Scipy
+:mod:`tensor.slinalg` --  Linear Algebra Ops Using Scipy (deprecated)
 ===================================================================
 
 .. module:: tensor.slinalg
    :platform: Unix, Windows
-   :synopsis: Linear Algebra Ops Using Scipy
+   :synopsis: Linear Algebra Ops Using Scipy (deprecated)
 .. moduleauthor:: LISA
 
-.. note::
-
-   This module is not imported by default. You need to import it to use it.
-
-API
-===
-
-.. automodule:: pytensor.tensor.slinalg
-    :members:
-
+.. deprecated:: 2.x
+   The ``slinalg`` module is deprecated. Use :mod:`pytensor.tensor.linalg` instead.
+   All symbols previously exported from ``slinalg`` are available from
+   ``pytensor.tensor.linalg`` — see :ref:`libdoc_linalg`. Imports from
+   ``slinalg`` will be removed in PyTensor 3.0.

pytensor/graph/rewriting/basic.py

@@ -1495,7 +1495,7 @@ def output_fn(fgraph, node, s):
         from pytensor.graph.rewriting.basic import PatternNodeRewriter
         from pytensor.graph.rewriting.unify import OpPattern, LiteralString
         from pytensor.tensor.blockwise import Blockwise
-        from pytensor.tensor.slinalg import Solve
+        from pytensor.tensor.linalg.solvers.general import Solve
 
         PatternNodeRewriter(
             (

pytensor/graph/rewriting/unify.py

@@ -287,7 +287,7 @@ class OpPattern:
         from pytensor.graph.rewriting.unify import OpPattern
         from pytensor.tensor.elemwise import CAReduce
         from pytensor.tensor.blockwise import Blockwise
-        from pytensor.tensor.slinalg import Solve
+        from pytensor.tensor.linalg.solvers.general import Solve
 
         @node_rewriter(tracks=[OpPattern(CAReduce, axis=None)])
         def local_car_reduce_all_rewriter(fgraph, node):
@@ -352,7 +352,7 @@ def output_fn(fgraph, node, s):
         import pytensor.tensor as pt
         from pytensor.graph.rewriting.unify import OpPattern
         from pytensor.tensor.blockwise import Blockwise
-        from pytensor.tensor.slinalg import Solve
+        from pytensor.tensor.linalg.solvers.general import Solve
 
         A = var("A")
         b = var("b")

pytensor/link/jax/dispatch/__init__.py

@@ -9,13 +9,12 @@
 import pytensor.link.jax.dispatch.extra_ops
 import pytensor.link.jax.dispatch.pad
 import pytensor.link.jax.dispatch.math
-import pytensor.link.jax.dispatch.nlinalg
+import pytensor.link.jax.dispatch.linalg
 import pytensor.link.jax.dispatch.random
 import pytensor.link.jax.dispatch.scalar
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.shape
 import pytensor.link.jax.dispatch.signal
-import pytensor.link.jax.dispatch.slinalg
 import pytensor.link.jax.dispatch.sort
 import pytensor.link.jax.dispatch.sparse
 import pytensor.link.jax.dispatch.subtensor

pytensor/link/jax/dispatch/linalg/__init__.py

@@ -0,0 +1,8 @@
+from pytensor.link.jax.dispatch.linalg import (
+    constructors,
+    decomposition,
+    inverse,
+    products,
+    solvers,
+    summary,
+)

pytensor/link/jax/dispatch/linalg/constructors.py

@@ -0,0 +1,12 @@
+import jax
+
+from pytensor.link.jax.dispatch.basic import jax_funcify
+from pytensor.tensor.linalg.constructors import BlockDiagonal
+
+
+@jax_funcify.register(BlockDiagonal)
+def jax_funcify_BlockDiagonalMatrix(op, **kwargs):
+    def block_diag(*inputs):
+        return jax.scipy.linalg.block_diag(*inputs)
+
+    return block_diag