Run doctest-modules on CI

.github/workflows/test.yml

@@ -78,6 +78,7 @@ jobs:
         install-jax: [0]
         install-torch: [0]
         part:
+          - "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
           - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
           - "tests/scan"
           - "tests/sparse"
@@ -96,6 +97,10 @@ jobs:
             part: "tests/tensor/test_math.py"
           - fast-compile: 1
             float32: 1
+          - part: "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
+            float32: 1
+          - part: "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
+            fast-compile: 1
         include:
           - install-numba: 1
             python-version: "3.10"
@@ -149,11 +154,12 @@ jobs:
         shell: micromamba-shell {0}
         run: |
 
-          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
+          micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 -c pytorch -c nvidia; fi
-
+          pip install pytest-sphinx
+          
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'

environment.yml

@@ -34,6 +34,8 @@ dependencies:
   - pytest-xdist
   - pytest-benchmark
   - pytest-mock
+  - pip:
+    - pytest-sphinx
   # For building docs
   - sphinx>=5.1.0,<6
   - sphinx_rtd_theme

pyproject.toml

@@ -75,6 +75,7 @@ tests = [
     "coverage>=5.1",
     "pytest-benchmark",
     "pytest-mock",
+    "pytest-sphinx",
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
 jax = ["jax", "jaxlib"]

pytensor/gradient.py

@@ -2243,7 +2243,7 @@ def grad_clip(x, lower_bound, upper_bound):
     >>> z2 = pytensor.gradient.grad(x**2, x)
     >>> f = pytensor.function([x], outputs = [z, z2])
     >>> print(f(2.0))
-    [array(1.0), array(4.0)]
+    [array(1.), array(4.)]
 
     Notes
     -----

pytensor/graph/basic.py

@@ -1034,7 +1034,10 @@ def orphans_between(
 
     Examples
     --------
-    >>> orphans_between([x], [(x+y).out])
+    >>> from pytensor.graph.basic import orphans_between
+    >>> from pytensor.tensor import scalars
+    >>> x, y = scalars("xy")
+    >>> list(orphans_between([x], [(x+y)]))
     [y]
 
     """

pytensor/misc/pkl_utils.py

@@ -239,7 +239,7 @@ def dump(
     >>> foo_2 = pytensor.shared(1, name='foo')
     >>> with open('model.zip', 'wb') as f:
     ...     dump((foo_1, foo_2, np.array(2)), f)
-    >>> np.load('model.zip').keys()
+    >>> list(np.load('model.zip').keys())
     ['foo', 'foo_2', 'array_0', 'pkl']
     >>> np.load('model.zip')['foo']
     array(0)

pytensor/scalar/basic.py

@@ -208,6 +208,7 @@ class autocast_float_as:
 
     Examples
     --------
+    >>> from pytensor.tensor import fvector
     >>> with autocast_float_as('float32'):
     ...    assert (fvector() + 1.1).dtype == 'float32'  # temporary downcasting
     >>> assert (fvector() + 1.1).dtype == 'float64' # back to default behaviour

pytensor/sparse/basic.py

@@ -4316,23 +4316,26 @@ def block_diag(*matrices: TensorVariable, format: Literal["csc", "csr"] = "csc")
     --------
     Create a sparse block diagonal matrix from two sparse 2x2 matrices:
 
-    ..code-block:: python
+    .. testcode::
         import numpy as np
         from pytensor.sparse import block_diag
         from scipy.sparse import csr_matrix
 
         A = csr_matrix([[1, 2], [3, 4]])
         B = csr_matrix([[5, 6], [7, 8]])
-        result_sparse = block_diag(A, B, format='csr', name='X')
+        result_sparse = block_diag(A, B, format='csr')
 
         print(result_sparse)
-        >>>  SparseVariable{csr,int32}
-
         print(result_sparse.toarray().eval())
-        >>> array([[1, 2, 0, 0],
-        >>> [3, 4, 0, 0],
-        >>> [0, 0, 5, 6],
-        >>> [0, 0, 7, 8]])
+
+    .. testoutput::
+
+        SparseVariable{csr,int64}
+        [[1 2 0 0]
+         [3 4 0 0]
+         [0 0 5 6]
+         [0 0 7 8]]
+
     """
     if len(matrices) == 1:
         return matrices

pytensor/tensor/basic.py

@@ -1112,23 +1112,24 @@ def tril(m, k=0):
 
     Examples
     --------
-    >>> at.tril(np.arange(1,13).reshape(4,3), -1).eval()
+    >>> import pytensor.tensor as pt
+    >>> pt.tril(pt.arange(1,13).reshape((4,3)), -1).eval()
     array([[ 0,  0,  0],
            [ 4,  0,  0],
            [ 7,  8,  0],
            [10, 11, 12]])
 
-    >>> at.tril(np.arange(3*4*5).reshape(3, 4, 5)).eval()
+    >>> pt.tril(pt.arange(3*4*5).reshape((3, 4, 5))).eval()
     array([[[ 0,  0,  0,  0,  0],
             [ 5,  6,  0,  0,  0],
             [10, 11, 12,  0,  0],
             [15, 16, 17, 18,  0]],
-
+    <BLANKLINE>
            [[20,  0,  0,  0,  0],
             [25, 26,  0,  0,  0],
             [30, 31, 32,  0,  0],
             [35, 36, 37, 38,  0]],
-
+    <BLANKLINE>
            [[40,  0,  0,  0,  0],
             [45, 46,  0,  0,  0],
             [50, 51, 52,  0,  0],
@@ -1154,23 +1155,24 @@ def triu(m, k=0):
 
     Examples
     --------
-    >>> at.triu(np.arange(1,13).reshape(4,3), -1).eval()
+    >>> import pytensor.tensor as pt
+    >>> pt.triu(pt.arange(1, 13).reshape((4, 3)), -1).eval()
     array([[ 1,  2,  3],
            [ 4,  5,  6],
            [ 0,  8,  9],
            [ 0,  0, 12]])
 
-    >>> at.triu(np.arange(3*4*5).reshape(3, 4, 5)).eval()
+    >>> pt.triu(np.arange(3*4*5).reshape((3, 4, 5))).eval()
     array([[[ 0,  1,  2,  3,  4],
             [ 0,  6,  7,  8,  9],
             [ 0,  0, 12, 13, 14],
             [ 0,  0,  0, 18, 19]],
-
+    <BLANKLINE>
            [[20, 21, 22, 23, 24],
             [ 0, 26, 27, 28, 29],
             [ 0,  0, 32, 33, 34],
             [ 0,  0,  0, 38, 39]],
-
+    <BLANKLINE>
            [[40, 41, 42, 43, 44],
             [ 0, 46, 47, 48, 49],
             [ 0,  0, 52, 53, 54],
@@ -2029,28 +2031,14 @@ def matrix_transpose(x: "TensorLike") -> TensorVariable:
 
     Examples
     --------
-    >>> import pytensor as pt
-    >>> import numpy as np
-    >>> x = np.arange(24).reshape((2, 3, 4))
-    [[[ 0  1  2  3]
-      [ 4  5  6  7]
-      [ 8  9 10 11]]
-
-     [[12 13 14 15]
-      [16 17 18 19]
-      [20 21 22 23]]]
+    >>> import pytensor.tensor as pt
+    >>> x = pt.arange(24).reshape((2, 3, 4))
+    >>> x.type.shape
+    (2, 3, 4)
 
+    >>> pt.matrix_transpose(x).type.shape
+    (2, 4, 3)
 
-    >>> pt.matrix_transpose(x).eval()
-    [[[ 0  4  8]
-      [ 1  5  9]
-      [ 2  6 10]
-      [ 3  7 11]]
-
-     [[12 16 20]
-      [13 17 21]
-      [14 18 22]
-      [15 19 23]]]
 
 
     Notes
@@ -2077,15 +2065,21 @@ class Split(COp):
 
     Examples
     --------
-    >>> x = vector()
-    >>> splits = lvector()
+    >>> from pytensor import function
+    >>> import pytensor.tensor as pt
+    >>> x = pt.vector(dtype="int")
+    >>> splits = pt.vector(dtype="int")
+
     You have to declare right away how many split_points there will be.
-    >>> ra, rb, rc = split(x, splits, n_splits = 3, axis = 0)
+    >>> ra, rb, rc = pt.split(x, splits, n_splits = 3, axis = 0)
     >>> f = function([x, splits], [ra, rb, rc])
     >>> a, b, c = f([0,1,2,3,4,5], [3, 2, 1])
-    a == [0,1,2]
-    b == [3, 4]
-    c == [5]
+    >>> a
+    array([0, 1, 2])
+    >>> b
+    array([3, 4])
+    >>> c
+    array([5])
 
     TODO: Don't make a copy in C impl
     """
@@ -2334,13 +2328,22 @@ class Join(COp):
 
     Examples
     --------
-    >>> x, y, z = tensor.matrix(), tensor.matrix(), tensor.matrix()
-    >>> u = tensor.vector()
+    >>> import pytensor.tensor as pt
+    >>> x, y, z = pt.matrix(), pt.matrix(), pt.matrix()
+    >>> u = pt.vector()
+
+    >>> r = pt.join(0, x, y, z)
+    >>> c = pt.join(1, x, y, z)
+
+    The axis has to be an index into the shape
+    >>> pt.join(2, x, y, z)
+    Traceback (most recent call last):
+    ValueError: Axis value 2 is out of range for the given input dimensions
 
-    >>> r = join(0, x, y, z)
-    >>> c = join(1, x, y, z)
-    >>> join(2, x, y, z)    # WRONG: the axis has to be an index into the shape
-    >>> join(0, x, u)       # WRONG: joined tensors must have the same rank
+    Joined tensors must have the same rank
+    >>> pt.join(0, x, u)
+    Traceback (most recent call last):
+    TypeError: Only tensors with the same number of dimensions can be joined. Input ndims were: [2, 1].
 
     """
 
@@ -3239,28 +3242,29 @@ class _nd_grid:
 
     Examples
     --------
-    >>> a = at.mgrid[0:5, 0:3]
+    >>> import pytensor.tensor as pt
+    >>> a = pt.mgrid[0:5, 0:3]
     >>> a[0].eval()
     array([[0, 0, 0],
            [1, 1, 1],
            [2, 2, 2],
            [3, 3, 3],
-           [4, 4, 4]], dtype=int8)
+           [4, 4, 4]])
     >>> a[1].eval()
     array([[0, 1, 2],
            [0, 1, 2],
            [0, 1, 2],
            [0, 1, 2],
-           [0, 1, 2]], dtype=int8)
-    >>> b = at.ogrid[0:5, 0:3]
+           [0, 1, 2]])
+    >>> b = pt.ogrid[0:5, 0:3]
     >>> b[0].eval()
     array([[0],
            [1],
            [2],
            [3],
-           [4]], dtype=int8)
+           [4]])
     >>> b[1].eval()
-    array([[0, 1, 2, 3]], dtype=int8)
+    array([[0, 1, 2]])
 
     """
 
@@ -3924,8 +3928,8 @@ def stacklists(arg):
     >>> X = stacklists([[a, b], [c, d]])
     >>> f = function([a, b, c, d], X)
     >>> f(1, 2, 3, 4)
-    array([[ 1.,  2.],
-           [ 3.,  4.]], dtype=float32)
+    array([[1., 2.],
+           [3., 4.]])
 
     We can also stack arbitrarily shaped tensors. Here we stack matrices into
     a 2 by 2 grid:

pytensor/tensor/extra_ops.py

@@ -254,7 +254,7 @@ def searchsorted(x, v, side="left", sorter=None):
     --------
     >>> from pytensor import tensor as pt
     >>> from pytensor.tensor import extra_ops
-    >>> x = ptb.dvector()
+    >>> x = pt.dvector("x")
     >>> idx = x.searchsorted(3)
     >>> idx.eval({x: [1,2,3,4,5]})
     array(2)
@@ -1165,12 +1165,12 @@ class Unique(Op):
     >>> x = pytensor.tensor.vector()
     >>> f = pytensor.function([x], Unique(True, True, False)(x))
     >>> f([1, 2., 3, 4, 3, 2, 1.])
-    [array([ 1.,  2.,  3.,  4.]), array([0, 1, 2, 3]), array([0, 1, 2, 3, 2, 1, 0])]
+    [array([1., 2., 3., 4.]), array([0, 1, 2, 3]), array([0, 1, 2, 3, 2, 1, 0])]
 
     >>> y = pytensor.tensor.matrix()
     >>> g = pytensor.function([y], Unique(True, True, False)(y))
     >>> g([[1, 1, 1.0], (2, 3, 3.0)])
-    [array([ 1.,  2.,  3.]), array([0, 3, 4]), array([0, 0, 0, 1, 2, 2])]
+    [array([1., 2., 3.]), array([0, 3, 4]), array([0, 0, 0, 1, 2, 2])]
 
     """