Standardize Op.grad arguments

Author: Armavica

pytensor/breakpoint.py

@@ -143,8 +143,8 @@ def perform(self, node, inputs, output_storage):
             for i in range(len(output_storage)):
                 output_storage[i][0] = inputs[i + 1]
 
-    def grad(self, inputs, output_gradients):
-        return [DisconnectedType()(), *output_gradients]
+    def grad(self, inputs, output_grads):
+        return [DisconnectedType()(), *output_grads]
 
     def infer_shape(self, fgraph, inputs, input_shapes):
         # Return the shape of every input but the condition (first input)

pytensor/compile/ops.py

@@ -92,8 +92,8 @@ def make_node(self, x):
     def infer_shape(self, fgraph, node, input_shapes):
         return input_shapes
 
-    def grad(self, args, g_outs):
-        return g_outs
+    def grad(self, inputs, output_grads):
+        return output_grads
 
 
 view_op = ViewOp()

pytensor/gradient.py

@@ -2316,8 +2316,8 @@ def _is_zero(x):
 
 
 class ZeroGrad(ViewOp):
-    def grad(self, args, g_outs):
-        return [g_out.zeros_like() for g_out in g_outs]
+    def grad(self, inputs, output_grads):
+        return [g_out.zeros_like() for g_out in output_grads]
 
     def R_op(self, inputs, eval_points):
         if eval_points[0] is None:
@@ -2354,8 +2354,8 @@ def zero_grad(x):
 
 
 class UndefinedGrad(ViewOp):
-    def grad(self, args, g_outs):
-        return [grad_undefined(self, i, arg) for i, arg in enumerate(args)]
+    def grad(self, inputs, output_grads):
+        return [grad_undefined(self, i, arg) for i, arg in enumerate(inputs)]
 
     def R_op(self, inputs, eval_points):
         return [None]
@@ -2392,8 +2392,8 @@ def undefined_grad(x):
 
 
 class DisconnectedGrad(ViewOp):
-    def grad(self, args, g_outs):
-        return [disconnected_type() for g_out in g_outs]
+    def grad(self, inputs, output_grads):
+        return [disconnected_type() for g_out in output_grads]
 
     def R_op(self, inputs, eval_points):
         return [None]
@@ -2447,10 +2447,10 @@ def __init__(self, clip_lower_bound, clip_upper_bound):
         if not self.clip_upper_bound >= self.clip_lower_bound:
             raise ValueError("`clip_upper_bound` should be >= `clip_lower_bound`")
 
-    def grad(self, args, g_outs):
+    def grad(self, inputs, output_grads):
         return [
             pytensor.tensor.clip(g_out, self.clip_lower_bound, self.clip_upper_bound)
-            for g_out in g_outs
+            for g_out in output_grads
         ]
 
 
@@ -2490,8 +2490,8 @@ class GradScale(ViewOp):
     def __init__(self, multiplier):
         self.multiplier = multiplier
 
-    def grad(self, args, g_outs):
-        return [self.multiplier * g_out for g_out in g_outs]
+    def grad(self, inputs, output_grads):
+        return [self.multiplier * g_out for g_out in output_grads]
 
 
 def grad_scale(x, multiplier):

pytensor/ifelse.py

@@ -239,10 +239,10 @@ def make_node(self, condition: "TensorLike", *true_false_branches: Any):
     def R_op(self, inputs, eval_points):
         return self(inputs[0], *eval_points[1:], return_list=True)
 
-    def grad(self, ins, grads):
-        condition = ins[0]
-        inputs_true_branch = ins[1:][: self.n_outs]
-        inputs_false_branch = ins[1:][self.n_outs :]
+    def grad(self, inputs, output_grads):
+        condition = inputs[0]
+        inputs_true_branch = inputs[1:][: self.n_outs]
+        inputs_false_branch = inputs[1:][self.n_outs :]
 
         if self.name is not None:
             nw_name_t = self.name + "_grad_t"
@@ -260,19 +260,19 @@ def grad(self, ins, grads):
         # dtypes.
         inputs_true_grad = (
             [condition]
-            + grads
+            + output_grads
             + [
-                pt.basic.zeros_like(t, dtype=grads[i].dtype)
+                pt.basic.zeros_like(t, dtype=output_grads[i].dtype)
                 for i, t in enumerate(inputs_true_branch)
             ]
         )
         inputs_false_grad = (
             [condition]
             + [
-                pt.basic.zeros_like(f, dtype=grads[i].dtype)
+                pt.basic.zeros_like(f, dtype=output_grads[i].dtype)
                 for i, f in enumerate(inputs_false_branch)
             ]
-            + grads
+            + output_grads
         )
 
         # `condition` does affect the elements of the output so it is connected.

pytensor/link/jax/ops.py

@@ -135,14 +135,14 @@ def perform_jax(self, *inputs):
             return outputs[0]
         return outputs
 
-    def grad(self, inputs, output_gradients):
+    def grad(self, inputs, output_grads):
         """Compute gradients using JAX's vector-Jacobian product (VJP)."""
         import jax
 
         # Find indices of outputs that need gradients
         connected_output_indices = [
             i
-            for i, output_grad in enumerate(output_gradients)
+            for i, output_grad in enumerate(output_grads)
             if not isinstance(output_grad.type, DisconnectedType)
         ]
 
@@ -190,7 +190,7 @@ def restricted_function(*input_values):
         )
 
         return vjp_op(
-            *[*inputs, *[output_gradients[i] for i in connected_output_indices]],
+            *[*inputs, *[output_grads[i] for i in connected_output_indices]],
             return_list=True,
         )
 

pytensor/printing.py

@@ -804,8 +804,8 @@ def perform(self, node, inputs, output_storage):
         xout[0] = xin
         self.global_fn(self, xin)
 
-    def grad(self, input, output_gradients):
-        return output_gradients
+    def grad(self, inputs, output_grads):
+        return output_grads
 
     def R_op(self, inputs, eval_points):
         return list(eval_points)

pytensor/raise_op.py

@@ -88,8 +88,8 @@ def perform(self, node, inputs, output_storage):
         if not all(conds):
             raise self.exc_type(self.msg)
 
-    def grad(self, input, output_gradients):
-        return output_gradients + [DisconnectedType()()] * (len(input) - 1)
+    def grad(self, inputs, output_grads):
+        return output_grads + [DisconnectedType()()] * (len(inputs) - 1)
 
     def connection_pattern(self, node):
         return [[1]] + [[0]] * (len(node.inputs) - 1)

pytensor/scalar/basic.py

@@ -1272,7 +1272,7 @@ def perform(self, node, inputs, output_storage):
     def impl(self, *inputs):
         raise MethodNotDefined("impl", type(self), self.__class__.__name__)
 
-    def grad(self, inputs, output_gradients):
+    def grad(self, inputs, output_grads):
         raise MethodNotDefined("grad", type(self), self.__class__.__name__)
 
     def L_op(self, inputs, outputs, output_grads):
@@ -1683,7 +1683,7 @@ def output_types(self, *input_types):
                 )
         return upcast_out(*input_types[0])
 
-    def grad(self, inputs, output_gradients):
+    def grad(self, inputs, output_grads):
         return [inputs[0].zeros_like(dtype=config.floatX)]
 
 
@@ -1701,7 +1701,7 @@ def output_types(self, *input_types):
                 )
         return upcast_out(*input_types[0])
 
-    def grad(self, inputs, output_gradients):
+    def grad(self, inputs, output_grads):
         a, b = inputs
         return [
             a.zeros_like(dtype=config.floatX),
@@ -1942,8 +1942,8 @@ def c_code(self, node, name, inputs, outputs, sub):
         else:
             return z + " = " + op.join(inputs) + ";"
 
-    def grad(self, inputs, gout):
-        (gz,) = gout
+    def grad(self, inputs, output_grads):
+        (gz,) = output_grads
         retval = []
 
         # The following 3 lines verify that gz is complex when the
@@ -2045,9 +2045,9 @@ def c_code(self, node, name, inputs, outputs, sub):
             return f"{z} = ((double){x}) / {y};"
         return f"{z} = {x} / {y};"
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x, y) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         if x.type in complex_types:
             raise NotImplementedError()
 
@@ -2166,7 +2166,7 @@ def c_code(self, node, name, inputs, outputs, sub):
     def c_code_cache_version(self):
         return (6,)
 
-    def grad(self, inputs, g_output):
+    def grad(self, inputs, output_grads):
         return [inp.zeros_like(dtype=config.floatX) for inp in inputs]
 
 
@@ -2440,9 +2440,9 @@ def connection_pattern(self, node):
 
         return [[False], [True]]
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (_x, y) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         if y.type in continuous_types:
             # x is disconnected because the elements of x are not used
             return DisconnectedType()(), gz
@@ -2466,9 +2466,9 @@ def c_code(self, node, name, inputs, outputs, sub):
         (z,) = outputs
         return f"{z} = {x};"
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         if x.type in continuous_types:
             return (gz,)
         else:
@@ -2505,9 +2505,9 @@ def c_code(self, node, name, inputs, outputs, sub):
             return f"{z} = ({x}) ? 1 : 0;"
         return f"{z} = ({node.outputs[0].type.dtype_specs()[1]}){x};"
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         if self.o_type in continuous_types:
             return [gz]
         else:
@@ -2636,9 +2636,9 @@ def impl(self, x):
         # casting to output type is handled by filter
         return np.sign(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         rval = x.zeros_like()
 
         if rval.type in discrete_types:
@@ -2677,9 +2677,9 @@ class Ceil(UnaryScalarOp):
     def impl(self, x):
         return np.ceil(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         rval = x.zeros_like()
 
         if rval.type in discrete_types:
@@ -2703,9 +2703,9 @@ class Floor(UnaryScalarOp):
     def impl(self, x):
         return np.floor(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         rval = x.zeros_like()
 
         if rval.type in discrete_types:
@@ -2729,9 +2729,9 @@ class Trunc(UnaryScalarOp):
     def impl(self, x):
         return np.trunc(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         return [x.zeros_like(dtype=config.floatX)]
 
     def c_code(self, node, name, inputs, outputs, sub):
@@ -2757,9 +2757,9 @@ class RoundHalfToEven(UnaryScalarOp):
     def impl(self, x):
         return np.round(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         rval = x.zeros_like()
 
         if rval.type in discrete_types:
@@ -2843,9 +2843,9 @@ class RoundHalfAwayFromZero(UnaryScalarOp):
     def impl(self, x):
         return round_half_away_from_zero_vec(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (_gz,) = gout
+        (_gz,) = output_grads
         rval = x.zeros_like()
 
         if rval.type in discrete_types:
@@ -3865,9 +3865,9 @@ class Real(UnaryScalarOp):
     def impl(self, x):
         return np.real(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (_x,) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         return [complex(gz, 0)]
 
     def c_code(self, *args, **kwargs):
@@ -3883,9 +3883,9 @@ class Imag(UnaryScalarOp):
     def impl(self, x):
         return np.imag(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x,) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         if x.type in complex_types:
             return [complex(0, gz)]
         elif x.type in float_types:
@@ -3906,7 +3906,7 @@ class Angle(UnaryScalarOp):
     def impl(self, x):
         return np.angle(x)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         # y = x.imag
         # r = sqrt(y**2 + x.real**2)
         # g = y/r
@@ -3918,7 +3918,7 @@ def grad(self, inputs, gout):
         #     theta = -numpy.arcsin(g)+numpy.pi
 
         (c,) = inputs
-        (gtheta,) = gout
+        (gtheta,) = output_grads
         x = real(c)
         y = imag(c)
         r = _abs(c)
@@ -3957,9 +3957,9 @@ def output_types_preference(x, y):
     def impl(self, x, y):
         return builtins.complex(x, y)
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (x, y) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         return [cast(real(gz), x.type.dtype), cast(imag(gz), y.type.dtype)]
 
     def c_code(self, *args, **kwargs):
@@ -4002,9 +4002,9 @@ def impl(self, r, theta):
         else:
             return np.complex128(builtins.complex(x, y))
 
-    def grad(self, inputs, gout):
+    def grad(self, inputs, output_grads):
         (r, theta) = inputs
-        (gz,) = gout
+        (gz,) = output_grads
         gr = gz * complex_from_polar(1, theta)
         gtheta = gz * complex_from_polar(r, -theta)
         return [gr, gtheta]