Remove depreciated AllocDiag Op, rename AllocDiag2 -> AllocDiag

Author: jessegrabowski

pytensor/link/jax/dispatch/tensor_basic.py

@@ -10,7 +10,7 @@
 from pytensor.tensor import get_vector_length
 from pytensor.tensor.basic import (
     Alloc,
-    AllocDiag2,
+    AllocDiag,
     AllocEmpty,
     ARange,
     ExtractDiag,
@@ -212,10 +212,10 @@ def tri(*args):
 
 
 @register_specialize
-@node_rewriter([AllocDiag2])
+@node_rewriter([AllocDiag])
 def eagerly_inline_alloc_diag(fgraph, node):
     """
-    Inline `AllocDiag2` OpFromGraph into the graph so the component Ops can themselves be jaxified
+    Inline `AllocDiag` OpFromGraph into the graph so the component Ops can themselves be jaxified
 
     Parameters
     ----------

pytensor/tensor/basic.py

@@ -3727,112 +3727,7 @@ def trace(a, offset=0, axis1=0, axis2=1):
     return diagonal(a, offset=offset, axis1=axis1, axis2=axis2).sum(-1)
 
 
-class AllocDiag(Op):
-    """An `Op` that copies a vector to the diagonal of a zero-ed matrix."""
-
-    __props__ = ("offset", "axis1", "axis2")
-
-    def __init__(self, offset=0, axis1=0, axis2=1):
-        """
-        Parameters
-        ----------
-        offset: int
-            Offset of the diagonal from the main diagonal defined by `axis1`
-            and `axis2`. Can be positive or negative.  Defaults to main
-            diagonal (i.e. 0).
-        axis1: int
-            Axis to be used as the first axis of the 2-D sub-arrays to which
-            the diagonals will be allocated.  Defaults to first axis (i.e. 0).
-        axis2: int
-            Axis to be used as the second axis of the 2-D sub-arrays to which
-            the diagonals will be allocated.  Defaults to second axis (i.e. 1).
-        """
-        warnings.warn(
-            "AllocDiag is deprecated. Use `alloc_diag` instead",
-            FutureWarning,
-        )
-        self.offset = offset
-        if axis1 < 0 or axis2 < 0:
-            raise NotImplementedError("AllocDiag does not support negative axis")
-        if axis1 == axis2:
-            raise ValueError("axis1 and axis2 cannot be the same")
-        self.axis1 = axis1
-        self.axis2 = axis2
-
-    def make_node(self, diag):
-        diag = as_tensor_variable(diag)
-        if diag.type.ndim < 1:
-            raise ValueError(
-                "AllocDiag needs an input with 1 or more dimensions", diag.type
-            )
-        return Apply(
-            self,
-            [diag],
-            [diag.type.clone(shape=(None,) * (diag.ndim + 1))()],
-        )
-
-    def perform(self, node, inputs, outputs):
-        (x,) = inputs
-        (z,) = outputs
-
-        axis1 = np.minimum(self.axis1, self.axis2)
-        axis2 = np.maximum(self.axis1, self.axis2)
-        offset = self.offset
-
-        # Create array with one extra dimension for resulting matrix
-        result_shape = x.shape[:-1] + (x.shape[-1] + abs(offset),) * 2
-        result = np.zeros(result_shape, dtype=x.dtype)
-
-        # Create slice for diagonal in final 2 axes
-        idxs = np.arange(x.shape[-1])
-        diagonal_slice = (len(result_shape) - 2) * [slice(None)] + [
-            idxs + np.maximum(0, -offset),
-            idxs + np.maximum(0, offset),
-        ]
-
-        # Fill in final 2 axes with x
-        result[tuple(diagonal_slice)] = x
-
-        if len(x.shape) > 1:
-            # Re-order axes so they correspond to diagonals at axis1, axis2
-            axes = list(range(len(x.shape[:-1])))
-            last_idx = axes[-1]
-            axes = axes[:axis1] + [last_idx + 1] + axes[axis1:]
-            axes = axes[:axis2] + [last_idx + 2] + axes[axis2:]
-            result = result.transpose(axes)
-
-        z[0] = result
-
-    def grad(self, inputs, gout):
-        (gz,) = gout
-        return [diagonal(gz, offset=self.offset, axis1=self.axis1, axis2=self.axis2)]
-
-    def infer_shape(self, fgraph, nodes, shapes):
-        (x_shape,) = shapes
-        axis1 = np.minimum(self.axis1, self.axis2)
-        axis2 = np.maximum(self.axis1, self.axis2)
-
-        result_shape = list(x_shape[:-1])
-        diag_shape = x_shape[-1] + abs(self.offset)
-        result_shape = result_shape[:axis1] + [diag_shape] + result_shape[axis1:]
-        result_shape = result_shape[:axis2] + [diag_shape] + result_shape[axis2:]
-        return [tuple(result_shape)]
-
-    def __setstate__(self, state):
-        if "view_map" in state:
-            del state["view_map"]
-
-        self.__dict__.update(state)
-
-        if "offset" not in state:
-            self.offset = 0
-        if "axis1" not in state:
-            self.axis1 = 0
-        if "axis2" not in state:
-            self.axis2 = 1
-
-
-class AllocDiag2(OpFromGraph):
+class AllocDiag(OpFromGraph):
     """
     Wrapper Op for alloc_diag graphs
     """
@@ -3883,7 +3778,7 @@ def alloc_diag(diag, offset=0, axis1=0, axis2=1):
         axes = axes[:axis2] + [last_idx + 2] + axes[axis2:]
         result = result.transpose(axes)
 
-    return AllocDiag2(
+    return AllocDiag(
         inputs=[diag], outputs=[result], offset=offset, axis1=axis1, axis2=axis2
     )(diag)
 

pytensor/tensor/rewriting/linalg.py

@@ -10,7 +10,7 @@
 )
 from pytensor.scalar.basic import Mul
 from pytensor.tensor.basic import (
-    AllocDiag2,
+    AllocDiag,
     Eye,
     TensorVariable,
     diagonal,
@@ -475,7 +475,7 @@ def rewrite_det_diag_to_prod_diag(fgraph, node):
     # Check for use of pt.diag first
     if (
         inputs.owner
-        and isinstance(inputs.owner.op, AllocDiag2)
+        and isinstance(inputs.owner.op, AllocDiag)
         and inputs.owner.op.offset == 0
     ):
         diag_input = inputs.owner.inputs[0]