Do not redefine DisconnectedType everytime

Author: ricardoV94

doc/extending/creating_an_op.rst

@@ -495,7 +495,7 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
 
     from pytensor.graph.op import Op
     from pytensor.graph.basic import Apply
-    from pytensor.gradient import DisconnectedType
+    from pytensor.gradient import DisconnectedType, disconnected_type
 
     class TransposeAndSumOp(Op):
         __props__ = ()
@@ -539,13 +539,13 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
             out1_grad, out2_grad = output_grads
 
             if isinstance(out1_grad.type, DisconnectedType):
-                x_grad = DisconnectedType()()
+                x_grad = disconnected_type()
             else:
                 # Transpose the last two dimensions of the output gradient
                 x_grad = pt.swapaxes(out1_grad, -1, -2)
 
             if isinstance(out2_grad.type, DisconnectedType):
-                y_grad = DisconnectedType()()
+                y_grad = disconnected_type()
             else:
                 # Broadcast the output gradient to the same shape as y
                 y_grad = pt.broadcast_to(pt.expand_dims(out2_grad, -1), y.shape)

pytensor/breakpoint.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import disconnected_type
 from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.op import Op
 from pytensor.tensor.basic import as_tensor_variable
@@ -142,7 +142,7 @@ def perform(self, node, inputs, output_storage):
                 output_storage[i][0] = inputs[i + 1]
 
     def grad(self, inputs, output_gradients):
-        return [DisconnectedType()(), *output_gradients]
+        return [disconnected_type(), *output_gradients]
 
     def infer_shape(self, fgraph, inputs, input_shapes):
         # Return the shape of every input but the condition (first input)

pytensor/raise_op.py

@@ -2,7 +2,7 @@
 
 from textwrap import indent
 
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import disconnected_type
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.replace import _vectorize_node
 from pytensor.link.c.op import COp
@@ -89,7 +89,10 @@ def perform(self, node, inputs, outputs):
             raise self.exc_type(self.msg)
 
     def grad(self, input, output_gradients):
-        return output_gradients + [DisconnectedType()()] * (len(input) - 1)
+        return [
+            *output_gradients,
+            *(disconnected_type() for _ in range(len(input) - 1)),
+        ]
 
     def connection_pattern(self, node):
         return [[1]] + [[0]] * (len(node.inputs) - 1)

pytensor/scalar/basic.py

@@ -22,7 +22,7 @@
 import pytensor
 from pytensor import printing
 from pytensor.configdefaults import config
-from pytensor.gradient import DisconnectedType, grad_undefined
+from pytensor.gradient import disconnected_type, grad_undefined
 from pytensor.graph.basic import Apply, Constant, Variable, clone
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import HasInnerGraph
@@ -2426,13 +2426,13 @@ def grad(self, inputs, gout):
         (gz,) = gout
         if y.type in continuous_types:
             # x is disconnected because the elements of x are not used
-            return DisconnectedType()(), gz
+            return disconnected_type(), gz
         else:
             # when y is discrete, we assume the function can be extended
             # to deal with real-valued inputs by rounding them to the
             # nearest integer. f(x+eps) thus equals f(x) so the gradient
             # is zero, not disconnected or undefined
-            return DisconnectedType()(), y.zeros_like(dtype=config.floatX)
+            return disconnected_type(), y.zeros_like(dtype=config.floatX)
 
 
 second = Second(name="second")

pytensor/scan/op.py

@@ -63,7 +63,14 @@
 from pytensor.compile.mode import Mode, get_mode
 from pytensor.compile.profiling import register_profiler_printer
 from pytensor.configdefaults import config
-from pytensor.gradient import DisconnectedType, NullType, Rop, grad, grad_undefined
+from pytensor.gradient import (
+    DisconnectedType,
+    NullType,
+    Rop,
+    disconnected_type,
+    grad,
+    grad_undefined,
+)
 from pytensor.graph.basic import (
     Apply,
     Variable,
@@ -3073,7 +3080,7 @@ def compute_all_gradients(known_grads):
         )
         outputs = local_op(*outer_inputs, return_list=True)
         # Re-order the gradients correctly
-        gradients = [DisconnectedType()()]
+        gradients = [disconnected_type()]  # n_steps is disconnected
 
         offset = info.n_mit_mot + info.n_mit_sot + info.n_sit_sot + n_sitsot_outs
         for p, (x, t) in enumerate(
@@ -3098,7 +3105,7 @@ def compute_all_gradients(known_grads):
                 else:
                     gradients.append(x[::-1])
             elif t == "disconnected":
-                gradients.append(DisconnectedType()())
+                gradients.append(disconnected_type())
             elif t == "through_untraced":
                 gradients.append(
                     grad_undefined(
@@ -3126,7 +3133,7 @@ def compute_all_gradients(known_grads):
                 else:
                     gradients.append(x[::-1])
             elif t == "disconnected":
-                gradients.append(DisconnectedType()())
+                gradients.append(disconnected_type())
             elif t == "through_untraced":
                 gradients.append(
                     grad_undefined(
@@ -3149,15 +3156,15 @@ def compute_all_gradients(known_grads):
                 if not isinstance(dC_dout.type, DisconnectedType) and connected:
                     disconnected = False
             if disconnected:
-                gradients.append(DisconnectedType()())
+                gradients.append(disconnected_type())
             else:
                 gradients.append(
                     grad_undefined(
                         self, idx, inputs[idx], "Shared Variable with update"
                     )
                 )
 
-        gradients += [DisconnectedType()() for _ in range(info.n_nit_sot)]
+        gradients.extend(disconnected_type() for _ in range(info.n_nit_sot))
         begin = end
 
         end = begin + n_sitsot_outs
@@ -3167,7 +3174,7 @@ def compute_all_gradients(known_grads):
             if t == "connected":
                 gradients.append(x[-1])
             elif t == "disconnected":
-                gradients.append(DisconnectedType()())
+                gradients.append(disconnected_type())
             elif t == "through_untraced":
                 gradients.append(
                     grad_undefined(
@@ -3195,7 +3202,7 @@ def compute_all_gradients(known_grads):
                 ):
                     disconnected = False
             if disconnected:
-                gradients[idx] = DisconnectedType()()
+                gradients[idx] = disconnected_type()
         return gradients
 
     def R_op(self, inputs, eval_points):

pytensor/sparse/basic.py

@@ -18,7 +18,7 @@
 from pytensor import _as_symbolic, as_symbolic
 from pytensor import scalar as ps
 from pytensor.configdefaults import config
-from pytensor.gradient import DisconnectedType, grad_undefined
+from pytensor.gradient import DisconnectedType, disconnected_type, grad_undefined
 from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.op import Op
 from pytensor.link.c.type import generic
@@ -480,9 +480,9 @@ def grad(self, inputs, gout):
         )
         return [
             g_data,
-            DisconnectedType()(),
-            DisconnectedType()(),
-            DisconnectedType()(),
+            disconnected_type(),
+            disconnected_type(),
+            disconnected_type(),
         ]
 
     def infer_shape(self, fgraph, node, shapes):
@@ -1940,7 +1940,7 @@ def grad(self, inputs, grads):
         gx = g_output
         gy = pytensor.tensor.subtensor.advanced_subtensor1(g_output, *idx_list)
 
-        return [gx, gy] + [DisconnectedType()()] * len(idx_list)
+        return [gx, gy, *(disconnected_type() for _ in range(len(idx_list)))]
 
 
 construct_sparse_from_list = ConstructSparseFromList()

pytensor/tensor/basic.py

@@ -22,7 +22,7 @@
 from pytensor import config, printing
 from pytensor import scalar as ps
 from pytensor.compile.builders import OpFromGraph
-from pytensor.gradient import DisconnectedType, grad_undefined
+from pytensor.gradient import DisconnectedType, disconnected_type, grad_undefined
 from pytensor.graph import RewriteDatabaseQuery
 from pytensor.graph.basic import Apply, Constant, Variable, equal_computations
 from pytensor.graph.fg import FunctionGraph, Output
@@ -1738,7 +1738,7 @@ def grad(self, inputs, grads):
         # the inputs that specify the shape. If you grow the
         # shape by epsilon, the existing elements do not
         # change.
-        return [gx] + [DisconnectedType()() for i in inputs[1:]]
+        return [gx, *(disconnected_type() for _ in range(len(inputs) - 1))]
 
     def R_op(self, inputs, eval_points):
         if eval_points[0] is None:
@@ -2277,7 +2277,7 @@ def L_op(self, inputs, outputs, g_outputs):
         return [
             join(axis, *new_g_outputs),
             grad_undefined(self, 1, axis),
-            DisconnectedType()(),
+            disconnected_type(),
         ]
 
     def R_op(self, inputs, eval_points):
@@ -3340,14 +3340,14 @@ def L_op(self, inputs, outputs, grads):
         if self.dtype in discrete_dtypes:
             return [
                 start.zeros_like(dtype=config.floatX),
-                DisconnectedType()(),
+                disconnected_type(),
                 step.zeros_like(dtype=config.floatX),
             ]
         else:
             num_steps_taken = outputs[0].shape[0]
             return [
                 gz.sum(),
-                DisconnectedType()(),
+                disconnected_type(),
                 (gz * arange(num_steps_taken, dtype=self.dtype)).sum(),
             ]
 
@@ -4374,7 +4374,7 @@ def connection_pattern(self, node):
         return [[False] for i in node.inputs]
 
     def grad(self, inputs, grads):
-        return [DisconnectedType()() for i in inputs]
+        return [disconnected_type() for _ in range(len(inputs))]
 
     def R_op(self, inputs, eval_points):
         return [zeros(inputs, self.dtype)]

pytensor/tensor/extra_ops.py

@@ -8,7 +8,6 @@
 import pytensor
 import pytensor.scalar.basic as ps
 from pytensor.gradient import (
-    DisconnectedType,
     _float_zeros_like,
     disconnected_type,
     grad_undefined,
@@ -716,7 +715,7 @@ def grad(self, inputs, gout):
         gx_transpose = ptb.zeros_like(x_transpose)[repeated_arange].inc(gz_transpose)
         gx = ptb.moveaxis(gx_transpose, 0, axis)
 
-        return [gx, DisconnectedType()()]
+        return [gx, disconnected_type()]
 
     def infer_shape(self, fgraph, node, ins_shapes):
         i0_shapes = ins_shapes[0]

pytensor/tensor/fft.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import disconnected_type
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.tensor.basic import as_tensor_variable
@@ -59,7 +59,7 @@ def grad(self, inputs, output_grads):
             + [slice(None)]
         )
         gout = set_subtensor(gout[idx], gout[idx] * 0.5)
-        return [irfft_op(gout, s), DisconnectedType()()]
+        return [irfft_op(gout, s), disconnected_type()]
 
     def connection_pattern(self, node):
         # Specify that shape input parameter has no connection to graph and gradients.
@@ -121,7 +121,7 @@ def grad(self, inputs, output_grads):
             + [slice(None)]
         )
         gf = set_subtensor(gf[idx], gf[idx] * 2)
-        return [gf, DisconnectedType()()]
+        return [gf, disconnected_type()]
 
     def connection_pattern(self, node):
         # Specify that shape input parameter has no connection to graph and gradients.

pytensor/tensor/nlinalg.py

@@ -8,7 +8,7 @@
 
 from pytensor import scalar as ps
 from pytensor.compile.builders import OpFromGraph
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import DisconnectedType, disconnected_type
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.tensor import TensorLike
@@ -652,8 +652,8 @@ def s_grad_only(
             ]
             if all(is_disconnected):
                 # This should never actually be reached by Pytensor -- the SVD Op should be pruned from the gradient
-                # graph if its fully disconnected. It is included for completeness.
-                return [DisconnectedType()()]  # pragma: no cover
+                # graph if it's fully disconnected. It is included for completeness.
+                return [disconnected_type()]  # pragma: no cover
 
             elif is_disconnected == [True, False, True]:
                 # This is the same as the compute_uv = False, so we can drop back to that simpler computation, without