Use dispatch for gufunc signature and (partially) implement Subtensor gufunc

Author: brandonwillard

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,
@@ -3469,6 +3470,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
 

aesara/tensor/blockwise.py

@@ -8,14 +8,14 @@
 from aesara.graph.null_type import NullType
 from aesara.graph.op import Op
 from aesara.scalar.basic import constant as scalar_constant
-from aesara.tensor import get_scalar_constant_value
-from aesara.tensor.basic import ExtractDiag, atleast_Nd
-from aesara.tensor.elemwise import DimShuffle, Elemwise
+from aesara.scalar.basic import int64
+from aesara.tensor import get_gufunc_signature, get_scalar_constant_value
+from aesara.tensor.basic import atleast_Nd
+from aesara.tensor.elemwise import DimShuffle
 from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.extra_ops import broadcast_shape
 from aesara.tensor.math import sum as at_sum
 from aesara.tensor.shape import shape_tuple
-from aesara.tensor.subtensor import Subtensor
 from aesara.tensor.type import TensorType
 
 
@@ -110,7 +110,15 @@ def get_dim_size(x):
         res = dim_sizes.get(x)
 
         if res is None:
-            return scalar_constant(int(x))
+            try:
+                return scalar_constant(int(x))
+            except (ValueError, TypeError):
+                # Return a symbolic placeholder for new dimension references
+                # For example, a signature like `("m", "n") -> ("p",)` means
+                # that there will be no `"p"` label to reference in `dim_sizes`
+                # (i.e. pre-existing dimension labels that already have values
+                # assigned to them).
+                return int64(name=x)
 
         return res
 
@@ -292,35 +300,12 @@ def transform(var: "TensorVariable", client_node: Optional[Apply]) -> Variable:
 
                 return atleast_Nd(res, n=nd)
 
-            if isinstance(node.op, (Subtensor, ExtractDiag)):
-                return var
-
             blocked_inputs = [transform(ipt, node) for ipt in node.inputs]
-            grad_signature = getattr(node.op, "gufunc_sig", None)
-            op = node.op
-
-            if grad_signature is None:
-                if isinstance(op, DimShuffle):
-                    # remove the extra dimensions that
-                    # we have added during op creation
-                    new_order = [i for i in op.new_order if i != "x"]
-
-                    # derive gufunc signature for DimShuffle
-                    input_signature: Tuple[str, ...] = tuple(
-                        f"a{i}" for i in range(len(new_order))
-                    )
-                    output_signature: Tuple[str, ...] = tuple(
-                        f"a{i}" if i != "x" else "1" for i in op.new_order
-                    )
-                    grad_signature = ((input_signature,), (output_signature,))
-                elif isinstance(op, Elemwise):
-                    op = op.scalar_op
-                    grad_signature = (((),) * len(blocked_inputs), ((),))
-                else:
-                    raise ValueError(f"'{op}' object has no attribute 'gufunc_sig'")
-
-            new_r = Blockwise(op, signature=grad_signature)(*blocked_inputs)
+            grad_signature = get_gufunc_signature(node.op, blocked_inputs)
+            new_r = Blockwise(node.op, signature=grad_signature)(*blocked_inputs)
+
             assert isinstance(new_r, Variable)
+
             return new_r
 
         ret = []

aesara/tensor/elemwise.py

@@ -19,7 +19,7 @@
 from aesara.scalar.basic import bool as scalar_bool
 from aesara.scalar.basic import identity as scalar_identity
 from aesara.scalar.basic import transfer_type, upcast
-from aesara.tensor import _get_vector_length, as_tensor_variable
+from aesara.tensor import _get_gufunc_signature, _get_vector_length, as_tensor_variable
 from aesara.tensor import elemwise_cgen as cgen
 from aesara.tensor import get_vector_length
 from aesara.tensor.type import (
@@ -275,6 +275,20 @@ def grad(self, inp, grads):
             ]
 
 
+@_get_gufunc_signature.register(DimShuffle)
+def _get_gufunc_signature_DimShuffle(op, blocked_inputs):
+    # remove the extra dimensions that
+    # we have added during op creation
+    new_order = [i for i in op.new_order if i != "x"]
+
+    # derive gufunc signature for DimShuffle
+    input_signature: Tuple[str, ...] = tuple(f"a{i}" for i in range(len(new_order)))
+    output_signature: Tuple[str, ...] = tuple(
+        f"a{i}" if i != "x" else "1" for i in op.new_order
+    )
+    return ((input_signature,), (output_signature,))
+
+
 class DimShufflePrinter(Printer):
     def __p(self, new_order, pstate, r):
         if new_order != () and new_order[0] == "x":
@@ -1222,6 +1236,12 @@ def c_code_cache_version_apply(self, node):
             return ()
 
 
+@_get_gufunc_signature.register(Elemwise)
+def _get_gufunc_signature_Elemwise(op, blocked_inputs):
+    op = op.scalar_op
+    return (((),) * len(blocked_inputs), ((),))
+
+
 class CAReduce(COp):
     """Reduces a scalar operation along specified axes.
 

aesara/tensor/subtensor.py

@@ -19,7 +19,12 @@
 from aesara.misc.safe_asarray import _asarray
 from aesara.printing import Printer, pprint, set_precedence
 from aesara.scalar.basic import ScalarConstant
-from aesara.tensor import _get_vector_length, as_tensor_variable, get_vector_length
+from aesara.tensor import (
+    _get_gufunc_signature,
+    _get_vector_length,
+    as_tensor_variable,
+    get_vector_length,
+)
 from aesara.tensor.basic import alloc, get_scalar_constant_value
 from aesara.tensor.elemwise import DimShuffle
 from aesara.tensor.exceptions import (
@@ -1200,6 +1205,29 @@ def R_op(self, inputs, eval_points):
         return self(eval_points[0], *inputs[1:], return_list=True)
 
 
+@_get_gufunc_signature.register(Subtensor)
+def _get_gufunc_signature_Subtensor(op, blocked_inputs):
+    min_base_dims = len(op.idx_list)
+    index_input_types = get_slice_elements(
+        op.idx_list, lambda entry: isinstance(entry, Type)
+    )
+
+    indexed_input_sig = tuple(f"a{i}" for i in range(min_base_dims))
+    index_input_sig = tuple(
+        ("1",) if typ.ndim == 0 else tuple(f"b{i}{j}" for j in range(typ.ndim))
+        for i, typ in enumerate(index_input_types)
+    )
+
+    # TODO: Compute the number of output dimensions
+    out_ndim = 1
+    output_sig = tuple(f"d{i}" for i in range(out_ndim))
+
+    input_signature: Tuple[str, ...] = (indexed_input_sig,) + index_input_sig
+    output_signature: Tuple[str, ...] = (output_sig,)
+
+    return (input_signature, output_signature)
+
+
 class SubtensorPrinter(Printer):
     def process(self, r, pstate):
         return self._process(r.owner.op.idx_list, r.owner.inputs, pstate)