Enable ruff to format code in docstrings

pyproject.toml

@@ -124,6 +124,9 @@ testpaths = "tests/"
 line-length = 88
 exclude = ["doc/", "pytensor/_version.py"]
 
+[tool.ruff.format]
+docstring-code-format = true
+
 [tool.ruff.lint]
 select = ["C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC"]
 ignore = ["C408", "C901", "E501", "E741", "RUF012", "PERF203", "ISC001"]

pytensor/compile/builders.py

@@ -190,7 +190,8 @@ class OpFromGraph(Op, HasInnerGraph):
 
         from pytensor import function, tensor as pt
         from pytensor.compile.builders import OpFromGraph
-        x, y, z = pt.scalars('xyz')
+
+        x, y, z = pt.scalars("xyz")
         e = x + y * z
         op = OpFromGraph([x, y, z], [e])
         # op behaves like a normal pytensor op
@@ -206,7 +207,7 @@ class OpFromGraph(Op, HasInnerGraph):
         from pytensor import config, function, tensor as pt
         from pytensor.compile.builders import OpFromGraph
 
-        x, y, z = pt.scalars('xyz')
+        x, y, z = pt.scalars("xyz")
         s = pytensor.shared(np.random.random((2, 2)).astype(config.floatX))
         e = x + y * z + s
         op = OpFromGraph([x, y, z], [e])
@@ -221,12 +222,16 @@ class OpFromGraph(Op, HasInnerGraph):
         from pytensor import function, tensor as pt, grad
         from pytensor.compile.builders import OpFromGraph
 
-        x, y, z = pt.scalars('xyz')
+        x, y, z = pt.scalars("xyz")
         e = x + y * z
+
+
         def rescale_dy(inps, outputs, out_grads):
             x, y, z = inps
-            g, = out_grads
-            return z*2
+            (g,) = out_grads
+            return z * 2
+
+
         op = OpFromGraph(
             [x, y, z],
             [e],
@@ -236,7 +241,7 @@ class OpFromGraph(Op, HasInnerGraph):
         dx, dy, dz = grad(e2, [x, y, z])
         fn = function([x, y, z], [dx, dy, dz])
         # the gradient wrt y is now doubled
-        fn(2., 3., 4.) # [1., 8., 3.]
+        fn(2.0, 3.0, 4.0)  # [1., 8., 3.]
 
     """
 

pytensor/gradient.py

@@ -692,25 +692,24 @@ def subgraph_grad(wrt, end, start=None, cost=None, details=False):
 
     .. code-block:: python
 
-        x, t = pytensor.tensor.fvector('x'), pytensor.tensor.fvector('t')
-        w1 = pytensor.shared(np.random.standard_normal((3,4)))
-        w2 = pytensor.shared(np.random.standard_normal((4,2)))
-        a1 = pytensor.tensor.tanh(pytensor.tensor.dot(x,w1))
-        a2 = pytensor.tensor.tanh(pytensor.tensor.dot(a1,w2))
+        x, t = pytensor.tensor.fvector("x"), pytensor.tensor.fvector("t")
+        w1 = pytensor.shared(np.random.standard_normal((3, 4)))
+        w2 = pytensor.shared(np.random.standard_normal((4, 2)))
+        a1 = pytensor.tensor.tanh(pytensor.tensor.dot(x, w1))
+        a2 = pytensor.tensor.tanh(pytensor.tensor.dot(a1, w2))
         cost2 = pytensor.tensor.sqr(a2 - t).sum()
         cost2 += pytensor.tensor.sqr(w2.sum())
         cost1 = pytensor.tensor.sqr(w1.sum())
 
-        params = [[w2],[w1]]
-        costs = [cost2,cost1]
+        params = [[w2], [w1]]
+        costs = [cost2, cost1]
         grad_ends = [[a1], [x]]
 
         next_grad = None
         param_grads = []
         for i in range(2):
             param_grad, next_grad = pytensor.subgraph_grad(
-                wrt=params[i], end=grad_ends[i],
-                start=next_grad, cost=costs[i]
+                wrt=params[i], end=grad_ends[i], start=next_grad, cost=costs[i]
             )
             next_grad = dict(zip(grad_ends[i], next_grad))
             param_grads.extend(param_grad)
@@ -1704,9 +1703,11 @@ def verify_grad(
 
     Examples
     --------
-    >>> verify_grad(pytensor.tensor.tanh,
-    ...             (np.asarray([[2, 3, 4], [-1, 3.3, 9.9]]),),
-    ...             rng=np.random.default_rng(23098))
+    >>> verify_grad(
+    ...     pytensor.tensor.tanh,
+    ...     (np.asarray([[2, 3, 4], [-1, 3.3, 9.9]]),),
+    ...     rng=np.random.default_rng(23098),
+    ... )
 
     Parameters
     ----------
@@ -2342,9 +2343,9 @@ def grad_clip(x, lower_bound, upper_bound):
     Examples
     --------
     >>> x = pytensor.tensor.type.scalar()
-    >>> z = pytensor.gradient.grad(grad_clip(x, -1, 1)**2, x)
+    >>> z = pytensor.gradient.grad(grad_clip(x, -1, 1) ** 2, x)
     >>> z2 = pytensor.gradient.grad(x**2, x)
-    >>> f = pytensor.function([x], outputs = [z, z2])
+    >>> f = pytensor.function([x], outputs=[z, z2])
     >>> print(f(2.0))
     [array(1.), array(4.)]
 
@@ -2383,7 +2384,7 @@ def grad_scale(x, multiplier):
     >>> fprime = pytensor.function([x], fp)
     >>> print(fprime(2))  # doctest: +ELLIPSIS
     -0.416...
-    >>> f_inverse=grad_scale(fx, -1.)
+    >>> f_inverse = grad_scale(fx, -1.0)
     >>> fpp = pytensor.grad(f_inverse, wrt=x)
     >>> fpprime = pytensor.function([x], fpp)
     >>> print(fpprime(2))  # doctest: +ELLIPSIS

pytensor/graph/basic.py

@@ -399,18 +399,24 @@ class Variable(Node, Generic[_TypeType, OptionalApplyType]):
         import pytensor
         import pytensor.tensor as pt
 
-        a = pt.constant(1.5)            # declare a symbolic constant
-        b = pt.fscalar()                # declare a symbolic floating-point scalar
+        a = pt.constant(1.5)  # declare a symbolic constant
+        b = pt.fscalar()  # declare a symbolic floating-point scalar
 
-        c = a + b                       # create a simple expression
+        c = a + b  # create a simple expression
 
-        f = pytensor.function([b], [c])   # this works because a has a value associated with it already
+        f = pytensor.function(
+            [b], [c]
+        )  # this works because a has a value associated with it already
 
-        assert 4.0 == f(2.5)            # bind 2.5 to an internal copy of b and evaluate an internal c
+        assert 4.0 == f(2.5)  # bind 2.5 to an internal copy of b and evaluate an internal c
 
-        pytensor.function([a], [c])       # compilation error because b (required by c) is undefined
+        pytensor.function(
+            [a], [c]
+        )  # compilation error because b (required by c) is undefined
 
-        pytensor.function([a,b], [c])     # compilation error because a is constant, it can't be an input
+        pytensor.function(
+            [a, b], [c]
+        )  # compilation error because a is constant, it can't be an input
 
 
     The python variables ``a, b, c`` all refer to instances of type
@@ -587,10 +593,10 @@ class Variable(Node, Generic[_TypeType, OptionalApplyType]):
 
         >>> import numpy as np
         >>> import pytensor.tensor as pt
-        >>> x = pt.dscalar('x')
-        >>> y = pt.dscalar('y')
+        >>> x = pt.dscalar("x")
+        >>> y = pt.dscalar("y")
         >>> z = x + y
-        >>> np.allclose(z.eval({x : 16.3, y : 12.1}), 28.4)
+        >>> np.allclose(z.eval({x: 16.3, y: 12.1}), 28.4)
         True
 
         We passed :meth:`eval` a dictionary mapping symbolic PyTensor
@@ -963,9 +969,9 @@ def explicit_graph_inputs(
         import pytensor.tensor as pt
         from pytensor.graph.basic import explicit_graph_inputs
 
-        x = pt.vector('x')
+        x = pt.vector("x")
         y = pt.constant(2)
-        z = pt.mul(x*y)
+        z = pt.mul(x * y)
 
         inputs = list(explicit_graph_inputs(z))
         f = pytensor.function(inputs, z)
@@ -1041,7 +1047,7 @@ def orphans_between(
     >>> from pytensor.graph.basic import orphans_between
     >>> from pytensor.tensor import scalars
     >>> x, y = scalars("xy")
-    >>> list(orphans_between([x], [(x+y)]))
+    >>> list(orphans_between([x], [(x + y)]))
     [y]
 
     """

pytensor/link/c/interface.py

@@ -30,7 +30,7 @@ class CLinkerObject:
         .. code-block:: python
 
             def c_headers(self, **kwargs):
-                return ['<iostream>', '<math.h>', '/full/path/to/header.h']
+                return ["<iostream>", "<math.h>", "/full/path/to/header.h"]
 
 
         """
@@ -54,7 +54,7 @@ class CLinkerObject:
         .. code-block:: python
 
             def c_header_dirs(self, **kwargs):
-                return ['/usr/local/include', '/opt/weirdpath/src/include']
+                return ["/usr/local/include", "/opt/weirdpath/src/include"]
 
         """
         return []
@@ -134,7 +134,7 @@ class CLinkerObject:
         .. code-block:: python
 
             def c_compile_args(self, **kwargs):
-                return ['-ffast-math']
+                return ["-ffast-math"]
 
         """
         return []

pytensor/link/c/params_type.py

@@ -29,7 +29,9 @@ In your Op sub-class:
 
 .. code-block:: python
 
-    params_type = ParamsType(attr1=TensorType('int32', shape=(None, None)), attr2=ScalarType('float64'))
+    params_type = ParamsType(
+        attr1=TensorType("int32", shape=(None, None)), attr2=ScalarType("float64")
+    )
 
 If your op contains attributes ``attr1`` **and** ``attr2``, the default ``op.get_params()``
 implementation will automatically try to look for it and generate an appropriate Params object.
@@ -77,26 +79,35 @@ enumerations will be directly available as ParamsType attributes.
     from pytensor.link.c.params_type import ParamsType
     from pytensor.link.c.type import EnumType, EnumList
 
-    wrapper = ParamsType(enum1=EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3'),
-                         enum2=EnumType(PI=3.14, EPSILON=0.001))
+    wrapper = ParamsType(
+        enum1=EnumList("CONSTANT_1", "CONSTANT_2", "CONSTANT_3"),
+        enum2=EnumType(PI=3.14, EPSILON=0.001),
+    )
 
     # Each enum constant is available as a wrapper attribute:
-    print(wrapper.CONSTANT_1, wrapper.CONSTANT_2, wrapper.CONSTANT_3,
-          wrapper.PI, wrapper.EPSILON)
+    print(
+        wrapper.CONSTANT_1,
+        wrapper.CONSTANT_2,
+        wrapper.CONSTANT_3,
+        wrapper.PI,
+        wrapper.EPSILON,
+    )
 
     # For convenience, you can also look for a constant by name with
     # ``ParamsType.get_enum()`` method.
-    pi = wrapper.get_enum('PI')
-    epsilon = wrapper.get_enum('EPSILON')
-    constant_2 = wrapper.get_enum('CONSTANT_2')
+    pi = wrapper.get_enum("PI")
+    epsilon = wrapper.get_enum("EPSILON")
+    constant_2 = wrapper.get_enum("CONSTANT_2")
     print(pi, epsilon, constant_2)
 
 This implies that a ParamsType cannot contain different enum types with common enum names::
 
     # Following line will raise an error,
     # as there is a "CONSTANT_1" defined both in enum1 and enum2.
-    wrapper = ParamsType(enum1=EnumList('CONSTANT_1', 'CONSTANT_2'),
-                         enum2=EnumType(CONSTANT_1=0, CONSTANT_3=5))
+    wrapper = ParamsType(
+        enum1=EnumList("CONSTANT_1", "CONSTANT_2"),
+        enum2=EnumType(CONSTANT_1=0, CONSTANT_3=5),
+    )
 
 If your enum types contain constant aliases, you can retrieve them from ParamsType
 with ``ParamsType.enum_from_alias(alias)`` method (see :class:`pytensor.link.c.type.EnumType`
@@ -104,11 +115,12 @@ for more info about enumeration aliases).
 
 .. code-block:: python
 
-    wrapper = ParamsType(enum1=EnumList('A', ('B', 'beta'), 'C'),
-                         enum2=EnumList(('D', 'delta'), 'E', 'F'))
+    wrapper = ParamsType(
+        enum1=EnumList("A", ("B", "beta"), "C"), enum2=EnumList(("D", "delta"), "E", "F")
+    )
     b1 = wrapper.B
-    b2 = wrapper.get_enum('B')
-    b3 = wrapper.enum_from_alias('beta')
+    b2 = wrapper.get_enum("B")
+    b3 = wrapper.enum_from_alias("beta")
     assert b1 == b2 == b3
 
 """
@@ -236,10 +248,13 @@ class Params(dict):
 
         from pytensor.link.c.params_type import ParamsType, Params
         from pytensor.scalar import ScalarType
+
         # You must create a ParamsType first:
-        params_type = ParamsType(attr1=ScalarType('int32'),
-                                 key2=ScalarType('float32'),
-                                 field3=ScalarType('int64'))
+        params_type = ParamsType(
+            attr1=ScalarType("int32"),
+            key2=ScalarType("float32"),
+            field3=ScalarType("int64"),
+        )
         # Then you can create a Params object with
         # the params type defined above and values for attributes.
         params = Params(params_type, attr1=1, key2=2.0, field3=3)
@@ -491,11 +506,13 @@ class ParamsType(CType):
             from pytensor.link.c.type import EnumType, EnumList
             from pytensor.scalar import ScalarType
 
-            wrapper = ParamsType(scalar=ScalarType('int32'),
-                                 letters=EnumType(A=1, B=2, C=3),
-                                 digits=EnumList('ZERO', 'ONE', 'TWO'))
-            print(wrapper.get_enum('C'))  # 3
-            print(wrapper.get_enum('TWO'))  # 2
+            wrapper = ParamsType(
+                scalar=ScalarType("int32"),
+                letters=EnumType(A=1, B=2, C=3),
+                digits=EnumList("ZERO", "ONE", "TWO"),
+            )
+            print(wrapper.get_enum("C"))  # 3
+            print(wrapper.get_enum("TWO"))  # 2
 
             # You can also directly do:
             print(wrapper.C)
@@ -520,17 +537,19 @@ class ParamsType(CType):
             from pytensor.link.c.type import EnumType, EnumList
             from pytensor.scalar import ScalarType
 
-            wrapper = ParamsType(scalar=ScalarType('int32'),
-                                 letters=EnumType(A=(1, 'alpha'), B=(2, 'beta'), C=3),
-                                 digits=EnumList(('ZERO', 'nothing'), ('ONE', 'unit'), ('TWO', 'couple')))
-            print(wrapper.get_enum('C'))  # 3
-            print(wrapper.get_enum('TWO'))  # 2
-            print(wrapper.enum_from_alias('alpha')) # 1
-            print(wrapper.enum_from_alias('nothing')) # 0
+            wrapper = ParamsType(
+                scalar=ScalarType("int32"),
+                letters=EnumType(A=(1, "alpha"), B=(2, "beta"), C=3),
+                digits=EnumList(("ZERO", "nothing"), ("ONE", "unit"), ("TWO", "couple")),
+            )
+            print(wrapper.get_enum("C"))  # 3
+            print(wrapper.get_enum("TWO"))  # 2
+            print(wrapper.enum_from_alias("alpha"))  # 1
+            print(wrapper.enum_from_alias("nothing"))  # 0
 
             # For the following, alias 'C' is not defined, so the method looks for
             # a constant named 'C', and finds it.
-            print(wrapper.enum_from_alias('C')) # 3
+            print(wrapper.enum_from_alias("C"))  # 3
 
         .. note::
 
@@ -567,12 +586,14 @@ class ParamsType(CType):
             from pytensor.tensor.type import dmatrix
             from pytensor.scalar import ScalarType
 
+
             class MyObject:
                 def __init__(self):
                     self.a = 10
                     self.b = numpy.asarray([[1, 2, 3], [4, 5, 6]])
 
-            params_type = ParamsType(a=ScalarType('int32'), b=dmatrix, c=ScalarType('bool'))
+
+            params_type = ParamsType(a=ScalarType("int32"), b=dmatrix, c=ScalarType("bool"))
 
             o = MyObject()
             value_for_c = False

pytensor/link/c/type.py

@@ -318,7 +318,7 @@ class EnumType(CType, dict):
     .. code-block:: python
 
         enum = EnumType(CONSTANT_1=1, CONSTANT_2=2.5, CONSTANT_3=False, CONSTANT_4=True)
-        print (enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4)
+        print(enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4)
         # will print 1 2.5 0 1
 
     In C code:
@@ -334,7 +334,7 @@ class EnumType(CType, dict):
 
     .. code-block:: python
 
-        enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype='size_t')
+        enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t")
         # In C code, the Op param will then be a ``size_t``.
 
     .. note::
@@ -349,8 +349,9 @@ class EnumType(CType, dict):
 
         .. code-block:: python
 
-            enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2,
-                            ctype='size_t', cname='MyEnumName')
+            enum = EnumType(
+                CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t", cname="MyEnumName"
+            )
 
     **Example with aliases**
 
@@ -359,7 +360,7 @@ class EnumType(CType, dict):
 
     To give an alias to a constant in the EnumType constructor, use the following key-value syntax::
 
-        constant_name=(constant_alias, constant_value)
+        constant_name = (constant_alias, constant_value)
 
     You can then retrieve a constant from an alias with method ``EnumType.fromalias()``.
 
@@ -372,23 +373,23 @@ class EnumType(CType, dict):
         from pytensor.link.c.type import EnumType
 
         # You can remark that constant 'C' does not have an alias.
-        enum = EnumType(A=('alpha', 1), B=('beta', 2), C=3, D=('delta', 4))
+        enum = EnumType(A=("alpha", 1), B=("beta", 2), C=3, D=("delta", 4))
 
         # Constants are all directly available by name.
         print(enum.A, enum.B, enum.C, enum.D)
 
         # But we can also now get some constants by alias.
-        a = enum.fromalias('alpha')
-        b = enum.fromalias('beta')
-        d = enum.fromalias('delta')
+        a = enum.fromalias("alpha")
+        b = enum.fromalias("beta")
+        d = enum.fromalias("delta")
 
         # If method fromalias() receives an unknown alias,
         # it will looks for a constant with this alias
         # as exact constant name.
-        c = enum.fromalias('C') # will get enum.C
+        c = enum.fromalias("C")  # will get enum.C
 
         # An alias defined in an EnumType will be correctly converted with non-strict filtering.
-        value = enum.filter('delta', strict=False)
+        value = enum.filter("delta", strict=False)
         # value now contains enum.D, ie. 4.
 
     .. note::
@@ -648,14 +649,24 @@ class EnumList(EnumType):
 
     Example::
 
-        enum = EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3', 'CONSTANT_4', 'CONSTANT_5')
-        print (enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4, enum.CONSTANT_5)
+        enum = EnumList(
+            "CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", "CONSTANT_5"
+        )
+        print(
+            enum.CONSTANT_1,
+            enum.CONSTANT_2,
+            enum.CONSTANT_3,
+            enum.CONSTANT_4,
+            enum.CONSTANT_5,
+        )
         # will print: 0 1 2 3 4
 
     Like :class:`EnumType`, you can also define the C type and a C name for the op param.
     Default C type is ``int``::
 
-        enum = EnumList('CONSTANT_1', 'CONSTANT_2', 'CONSTANT_3', 'CONSTANT_4', ctype='unsigned int')
+        enum = EnumList(
+            "CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", ctype="unsigned int"
+        )
 
     Like :class:`EnumType`, you can also add an alias to a constant, by replacing the only constant name
     (e.g. ``'CONSTANT_NAME'``) by a couple with constant name first and constant alias second
@@ -663,7 +674,7 @@ class EnumList(EnumType):
 
     .. code-block:: python
 
-        enum = EnumList(('A', 'alpha'), ('B', 'beta'), 'C', 'D', 'E', 'F', ('G', 'gamma'))
+        enum = EnumList(("A", "alpha"), ("B", "beta"), "C", "D", "E", "F", ("G", "gamma"))
 
     See test class :class:`tests.graph.test_types.TestOpEnumList` for a working example.
 
@@ -727,7 +738,9 @@ class CEnumType(EnumList):
 
     .. code-block:: python
 
-        enum = CEnumType('CONSTANT_CNAME_1', 'CONSTANT_CNAME_2', 'CONSTANT_CNAME_3', ctype='long')
+        enum = CEnumType(
+            "CONSTANT_CNAME_1", "CONSTANT_CNAME_2", "CONSTANT_CNAME_3", ctype="long"
+        )
 
     Like :class:`EnumList`, you can also add an alias to a constant, with same syntax as in :class:`EnumList`.
 

pytensor/link/numba/dispatch/elemwise.py

@@ -185,7 +185,9 @@ def create_axis_reducer(
     .. code-block:: python
 
         def careduce_axis(x):
-            res_shape = tuple(shape[i] if i < axis else shape[i + 1] for i in range(ndim - 1))
+            res_shape = tuple(
+                shape[i] if i < axis else shape[i + 1] for i in range(ndim - 1)
+            )
             res = np.full(res_shape, identity, dtype=dtype)
 
             x_axis_first = x.transpose(reaxis_first)

pytensor/printing.py

@@ -1247,10 +1247,11 @@ def pydotprint(
         .. code-block:: python
 
             import pytensor
+
             v = pytensor.tensor.vector()
             from IPython.display import SVG
-            SVG(pytensor.printing.pydotprint(v*2, return_image=True,
-                                           format='svg'))
+
+            SVG(pytensor.printing.pydotprint(v * 2, return_image=True, format="svg"))
 
     In the graph, ellipses are Apply Nodes (the execution of an op)
     and boxes are variables.  If variables have names they are used as

pytensor/scalar/basic.py

@@ -209,9 +209,9 @@ 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
+    >>> with autocast_float_as("float32"):
+    ...     assert (fvector() + 1.1).dtype == "float32"  # temporary downcasting
+    >>> assert (fvector() + 1.1).dtype == "float64"  # back to default behaviour
 
     """
 

pytensor/scan/basic.py

@@ -207,13 +207,20 @@ def scan(
 
         .. code-block:: python
 
-            scan(fn, sequences = [ dict(input= Sequence1, taps = [-3,2,-1])
-                                 , Sequence2
-                                 , dict(input =  Sequence3, taps = 3) ]
-                   , outputs_info = [ dict(initial =  Output1, taps = [-3,-5])
-                                    , dict(initial = Output2, taps = None)
-                                    , Output3 ]
-                   , non_sequences = [ Argument1, Argument2])
+            scan(
+                fn,
+                sequences=[
+                    dict(input=Sequence1, taps=[-3, 2, -1]),
+                    Sequence2,
+                    dict(input=Sequence3, taps=3),
+                ],
+                outputs_info=[
+                    dict(initial=Output1, taps=[-3, -5]),
+                    dict(initial=Output2, taps=None),
+                    Output3,
+                ],
+                non_sequences=[Argument1, Argument2],
+            )
 
         `fn` should expect the following arguments in this given order:
 
@@ -240,11 +247,12 @@ def scan(
 
             import pytensor.tensor as pt
 
-            W   = pt.matrix()
+            W = pt.matrix()
             W_2 = W**2
 
+
             def f(x):
-                return pt.dot(x,W_2)
+                return pt.dot(x, W_2)
 
         The function `fn` is expected to return two things. One is a list of
         outputs ordered in the same order as `outputs_info`, with the
@@ -266,7 +274,7 @@ def scan(
         .. code-block:: python
 
             ...
-            return [y1_t, y2_t], {x:x+1}, until(x < 50)
+            return [y1_t, y2_t], {x: x + 1}, until(x < 50)
 
         Note that a number of steps--considered in here as the maximum
         number of steps--is still required even though a condition is

pytensor/tensor/basic.py

@@ -1113,13 +1113,13 @@ def tril(m, k=0):
     Examples
     --------
     >>> import pytensor.tensor as pt
-    >>> pt.tril(pt.arange(1,13).reshape((4,3)), -1).eval()
+    >>> pt.tril(pt.arange(1, 13).reshape((4, 3)), -1).eval()
     array([[ 0,  0,  0],
            [ 4,  0,  0],
            [ 7,  8,  0],
            [10, 11, 12]])
 
-    >>> pt.tril(pt.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],
@@ -1162,7 +1162,7 @@ def triu(m, k=0):
            [ 0,  8,  9],
            [ 0,  0, 12]])
 
-    >>> pt.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],
@@ -2108,9 +2108,9 @@ class Split(COp):
     >>> splits = pt.vector(dtype="int")
 
     You have to declare right away how many split_points there will be.
-    >>> ra, rb, rc = pt.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, b, c = f([0, 1, 2, 3, 4, 5], [3, 2, 1])
     >>> a
     array([0, 1, 2])
     >>> b
@@ -2831,28 +2831,28 @@ def stack(tensors: Sequence["TensorLike"], axis: int = 0):
     >>> b = pytensor.tensor.type.scalar()
     >>> c = pytensor.tensor.type.scalar()
     >>> x = pytensor.tensor.stack([a, b, c])
-    >>> x.ndim # x is a vector of length 3.
+    >>> x.ndim  # x is a vector of length 3.
     1
     >>> a = pytensor.tensor.type.tensor4()
     >>> b = pytensor.tensor.type.tensor4()
     >>> c = pytensor.tensor.type.tensor4()
     >>> x = pytensor.tensor.stack([a, b, c])
-    >>> x.ndim # x is a 5d tensor.
+    >>> x.ndim  # x is a 5d tensor.
     5
     >>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
-    >>> rval.shape # 3 tensors are stacked on axis 0
+    >>> rval.shape  # 3 tensors are stacked on axis 0
     (3, 2, 2, 2, 2)
     >>> x = pytensor.tensor.stack([a, b, c], axis=3)
     >>> x.ndim
     5
     >>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
-    >>> rval.shape # 3 tensors are stacked on axis 3
+    >>> rval.shape  # 3 tensors are stacked on axis 3
     (2, 2, 2, 3, 2)
     >>> x = pytensor.tensor.stack([a, b, c], axis=-2)
     >>> x.ndim
     5
     >>> rval = x.eval(dict((t, np.zeros((2, 2, 2, 2))) for t in [a, b, c]))
-    >>> rval.shape # 3 tensors are stacked on axis -2
+    >>> rval.shape  # 3 tensors are stacked on axis -2
     (2, 2, 2, 3, 2)
     """
     if not isinstance(tensors, Sequence):
@@ -3892,7 +3892,7 @@ def stacklists(arg):
     >>> from pytensor.tensor import stacklists
     >>> from pytensor.tensor.type import scalars, matrices
     >>> from pytensor import function
-    >>> a, b, c, d = scalars('abcd')
+    >>> a, b, c, d = scalars("abcd")
     >>> X = stacklists([[a, b], [c, d]])
     >>> f = function([a, b, c, d], X)
     >>> f(1, 2, 3, 4)
@@ -3903,10 +3903,10 @@ def stacklists(arg):
     a 2 by 2 grid:
 
     >>> from numpy import ones
-    >>> a, b, c, d = matrices('abcd')
+    >>> a, b, c, d = matrices("abcd")
     >>> X = stacklists([[a, b], [c, d]])
     >>> f = function([a, b, c, d], X)
-    >>> x = ones((4, 4), 'float32')
+    >>> x = ones((4, 4), "float32")
     >>> f(x, x, x, x).shape
     (2, 2, 4, 4)
 

pytensor/tensor/einsum.py

@@ -467,6 +467,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
     .. code-block:: python
 
         import pytensor as pt
+
         A = pt.matrix("A")
         B = pt.matrix("B")
         C = pt.einsum("ij, jk -> ik", A, B)
@@ -481,6 +482,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
     .. code-block:: python
 
         import pytensor as pt
+
         A = pt.tensor("A", shape=(None, 4, 5))
         B = pt.tensor("B", shape=(None, 5, 6))
         C = pt.einsum("bij, bjk -> bik", A, B)
@@ -496,6 +498,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
     .. code-block:: python
 
         import pytensor as pt
+
         A = pt.tensor("A", shape=(4, None, None, None, 5))
         B = pt.tensor("B", shape=(5, None, None, None, 6))
         C = pt.einsum("i...j, j...k -> ...ik", A, B)
@@ -510,6 +513,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
     .. code-block:: python
 
         import pytensor as pt
+
         x = pt.tensor("x", shape=(3,))
         y = pt.tensor("y", shape=(4,))
         z = pt.einsum("i, j -> ij", x, y)

pytensor/tensor/elemwise.py

@@ -77,7 +77,7 @@ class DimShuffle(ExternalCOp):
 
     .. code-block:: python
 
-        DimShuffle((False, False, False), ['x', 2, 'x', 0, 1])
+        DimShuffle((False, False, False), ["x", 2, "x", 0, 1])
 
     This `Op` will only work on 3d tensors with no broadcastable
     dimensions.  The first dimension will be broadcastable,
@@ -101,16 +101,16 @@ class DimShuffle(ExternalCOp):
     --------
     .. code-block:: python
 
-        DimShuffle((), ['x'])  # make a 0d (scalar) into a 1d vector
+        DimShuffle((), ["x"])  # make a 0d (scalar) into a 1d vector
         DimShuffle((False, False), [0, 1])  # identity
         DimShuffle((False, False), [1, 0])  # inverts the 1st and 2nd dimensions
-        DimShuffle((False,), ['x', 0])  # make a row out of a 1d vector
-                                        # (N to 1xN)
-        DimShuffle((False,), [0, 'x'])  # make a column out of a 1d vector
-                                        # (N to Nx1)
+        DimShuffle((False,), ["x", 0])  # make a row out of a 1d vector
+        # (N to 1xN)
+        DimShuffle((False,), [0, "x"])  # make a column out of a 1d vector
+        # (N to Nx1)
         DimShuffle((False, False, False), [2, 0, 1])  # AxBxC to CxAxB
-        DimShuffle((False, False), [0, 'x', 1])  # AxB to Ax1xB
-        DimShuffle((False, False), [1, 'x', 0])  # AxB to Bx1xA
+        DimShuffle((False, False), [0, "x", 1])  # AxB to Ax1xB
+        DimShuffle((False, False), [1, "x", 0])  # AxB to Bx1xA
 
     The reordering of the dimensions can be done with the numpy.transpose
     function.
@@ -1180,15 +1180,15 @@ class CAReduce(COp):
     -----
     .. code-block:: python
 
-        CAReduce(add)      # sum (ie, acts like the numpy sum operation)
-        CAReduce(mul)      # product
+        CAReduce(add)  # sum (ie, acts like the numpy sum operation)
+        CAReduce(mul)  # product
         CAReduce(maximum)  # max
         CAReduce(minimum)  # min
-        CAReduce(or_)      # any # not lazy
-        CAReduce(and_)     # all # not lazy
-        CAReduce(xor)      # a bit at 1 tell that there was an odd number of
-                           # bit at that position that where 1. 0 it was an
-                           # even number ...
+        CAReduce(or_)  # any # not lazy
+        CAReduce(and_)  # all # not lazy
+        CAReduce(xor)  # a bit at 1 tell that there was an odd number of
+        # bit at that position that where 1. 0 it was an
+        # even number ...
 
     In order to (eventually) optimize memory usage patterns,
     `CAReduce` makes zero guarantees on the order in which it

pytensor/tensor/extra_ops.py

@@ -267,13 +267,13 @@ def searchsorted(x, v, side="left", sorter=None):
     >>> from pytensor.tensor import extra_ops
     >>> x = pt.dvector("x")
     >>> idx = x.searchsorted(3)
-    >>> idx.eval({x: [1,2,3,4,5]})
+    >>> idx.eval({x: [1, 2, 3, 4, 5]})
     array(2)
-    >>> extra_ops.searchsorted([1,2,3,4,5], 3).eval()
+    >>> extra_ops.searchsorted([1, 2, 3, 4, 5], 3).eval()
     array(2)
-    >>> extra_ops.searchsorted([1,2,3,4,5], 3, side='right').eval()
+    >>> extra_ops.searchsorted([1, 2, 3, 4, 5], 3, side="right").eval()
     array(3)
-    >>> extra_ops.searchsorted([1,2,3,4,5], [-10, 10, 2, 3]).eval()
+    >>> extra_ops.searchsorted([1, 2, 3, 4, 5], [-10, 10, 2, 3]).eval()
     array([0, 5, 1, 2])
 
     .. versionadded:: 0.9
@@ -1176,7 +1176,7 @@ class Unique(Op):
 
     >>> x = pytensor.tensor.vector()
     >>> f = pytensor.function([x], Unique(True, True, False)(x))
-    >>> f([1, 2., 3, 4, 3, 2, 1.])
+    >>> f([1, 2.0, 3, 4, 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()

pytensor/tensor/functional.py

@@ -39,9 +39,11 @@ def vectorize(func: Callable, signature: str | None = None) -> Callable:
         import pytensor
         import pytensor.tensor as pt
 
+
         def func(x):
             return pt.exp(x) / pt.sum(pt.exp(x))
 
+
         vec_func = pt.vectorize(func, signature="(a)->(a)")
 
         x = pt.matrix("x")
@@ -58,9 +60,11 @@ def vectorize(func: Callable, signature: str | None = None) -> Callable:
         import pytensor
         import pytensor.tensor as pt
 
+
         def func(x):
             return x[0], x[-1]
 
+
         vec_func = pt.vectorize(func, signature="(a)->(),()")
 
         x = pt.matrix("x")

pytensor/tensor/io.py

@@ -80,8 +80,8 @@ def load(path, dtype, shape, mmap_mode=None):
     --------
     >>> from pytensor import *
     >>> path = Variable(Generic(), None)
-    >>> x = tensor.load(path, 'int64', (None,))
-    >>> y = x*2
+    >>> x = tensor.load(path, "int64", (None,))
+    >>> y = x * 2
     >>> fn = function([path], y)
     >>> fn("stored-array.npy")  # doctest: +SKIP
     array([0, 2, 4, 6, 8], dtype=int64)

pytensor/tensor/math.py

@@ -2098,27 +2098,27 @@ def tensordot(
     are compatible. The resulting tensor will have shape (2, 5, 6) -- the
     dimensions that are not being summed:
 
-    >>> a = np.random.random((2,3,4))
-    >>> b = np.random.random((5,6,4,3))
+    >>> a = np.random.random((2, 3, 4))
+    >>> b = np.random.random((5, 6, 4, 3))
 
     #tensordot
-    >>> c = np.tensordot(a, b, [[1,2],[3,2]])
+    >>> c = np.tensordot(a, b, [[1, 2], [3, 2]])
 
     #loop replicating tensordot
     >>> a0, a1, a2 = a.shape
     >>> b0, b1, _, _ = b.shape
-    >>> cloop = np.zeros((a0,b0,b1))
+    >>> cloop = np.zeros((a0, b0, b1))
 
     #loop over non-summed indices -- these exist
     #in the tensor product.
     >>> for i in range(a0):
     ...     for j in range(b0):
     ...         for k in range(b1):
-    ...             #loop over summed indices -- these don't exist
-    ...             #in the tensor product.
+    ...             # loop over summed indices -- these don't exist
+    ...             # in the tensor product.
     ...             for l in range(a1):
     ...                 for m in range(a2):
-    ...                     cloop[i,j,k] += a[i,l,m] * b[j,k,m,l]
+    ...                     cloop[i, j, k] += a[i, l, m] * b[j, k, m, l]
 
     >>> np.allclose(c, cloop)
     True

pytensor/tensor/utils.py

@@ -61,8 +61,9 @@ def shape_of_variables(
     --------
     >>> import pytensor.tensor as pt
     >>> from pytensor.graph.fg import FunctionGraph
-    >>> x = pt.matrix('x')
-    >>> y = x[512:]; y.name = 'y'
+    >>> x = pt.matrix("x")
+    >>> y = x[512:]
+    >>> y.name = "y"
     >>> fgraph = FunctionGraph([x], [y], clone=False)
     >>> d = shape_of_variables(fgraph, {x: (1024, 1024)})
     >>> d[y]