Add Blockwise Op

aesara/tensor/__init__.py

@@ -98,6 +98,20 @@ def _get_vector_length_Constant(op: Union[Op, Variable], var: Constant) -> int:
     return len(var.data)
 
 
+def get_gufunc_signature(op, blocked_inputs):
+    sig = getattr(op, "gufunc_sig", None)
+
+    if sig is None:
+        return _get_gufunc_signature(op, blocked_inputs)
+
+    return sig
+
+
+@singledispatch
+def _get_gufunc_signature(op, blocked_inputs):
+    raise ValueError(f"'{op}' object has no attribute 'gufunc_sig'")
+
+
 import aesara.tensor.exceptions  # noqa
 from aesara.gradient import consider_constant, grad, hessian, jacobian  # noqa
 

aesara/tensor/basic.py

@@ -38,6 +38,7 @@
 from aesara.scalar.basic import ScalarConstant, ScalarVariable
 from aesara.tensor import (
     _as_tensor_variable,
+    _get_gufunc_signature,
     _get_vector_length,
     as_tensor_variable,
     get_vector_length,
@@ -986,7 +987,15 @@ def nonzero_values(a):
 
 class Tri(Op):
 
-    __props__ = ("dtype",)
+    gufunc_sig = (
+        (
+            (),
+            (),
+            (),
+        ),
+        (("n", "m"),),
+    )
+    __props__ = ("dtype", "gufunc_sig")
 
     def __init__(self, dtype=None):
         if dtype is None:
@@ -3470,6 +3479,12 @@ def __setstate__(self, state):
             self.axis2 = 1
 
 
+@_get_gufunc_signature.register(ExtractDiag)
+def _get_gufunc_signature_ExtractDiag(op, blocked_inputs):
+    # TODO:
+    raise NotImplementedError()
+
+
 extract_diag = ExtractDiag()
 # TODO: optimization to insert ExtractDiag with view=True
 
@@ -3502,7 +3517,8 @@ class AllocDiag(Op):
     It does the inverse of `ExtractDiag`.
     """
 
-    __props__ = ("offset", "axis1", "axis2")
+    gufunc_sig = (((),), (("m", "m"),))
+    __props__ = ("offset", "axis1", "axis2", "gufunc_sig")
 
     def __init__(self, offset=0, axis1=0, axis2=1):
         """