Remove AllocDiag2 from graphs in the specialization phase, after other rewrites that need it have fired.

Author: jessegrabowski

pytensor/link/jax/dispatch/tensor_basic.py

@@ -3,11 +3,14 @@
 import jax.numpy as jnp
 import numpy as np
 
+import pytensor
+from pytensor.graph import node_rewriter
 from pytensor.graph.basic import Constant
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor import get_vector_length
 from pytensor.tensor.basic import (
     Alloc,
+    AllocDiag2,
     AllocEmpty,
     ARange,
     ExtractDiag,
@@ -21,6 +24,7 @@
     get_underlying_scalar_constant_value,
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
+from pytensor.tensor.rewriting.basic import register_specialize
 from pytensor.tensor.shape import Shape_i
 
 
@@ -205,3 +209,28 @@ def tri(*args):
         return jnp.tri(*args, dtype=op.dtype)
 
     return tri
+
+
+@register_specialize
+@node_rewriter([AllocDiag2])
+def eagerly_inline_alloc_diag(fgraph, node):
+    """
+    Inline `AllocDiag2` OpFromGraph into the graph so the component Ops can themselves be jaxified
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        The function graph being rewritten
+    node: Apply
+        Node of the function graph to be optimized
+
+    Returns
+    -------
+
+    """
+    [input] = node.inputs
+    [output] = node.op.inner_outputs
+    inner_input = output.owner.inputs[1]
+
+    inline = pytensor.clone_replace(output, {inner_input: input})
+
+    return [inline]