Refactor get_scalar_constant_value

aesara/configdefaults.py

@@ -1462,13 +1462,6 @@ def add_scan_configvars():
         in_c_key=False,
     )
 
-    config.add(
-        "scan__debug",
-        "If True, enable extra verbose output related to scan",
-        BoolParam(False),
-        in_c_key=False,
-    )
-
 
 def _get_default_gpuarray__cache_path():
     return os.path.join(config.compiledir, "gpuarray_kernels")

aesara/graph/features.py

@@ -553,12 +553,9 @@ def replace_all_validate(
         self, fgraph, replacements, reason=None, verbose=None, **kwargs
     ):
         chk = fgraph.checkpoint()
+
         if verbose is None:
             verbose = config.optimizer_verbose
-        if config.scan__debug:
-            from aesara.scan.op import Scan
-
-            scans = [n for n in fgraph.apply_nodes if isinstance(n.op, Scan)]
 
         for r, new_r in replacements:
             try:
@@ -596,29 +593,14 @@ def replace_all_validate(
         except Exception as e:
             fgraph.revert(chk)
             if verbose:
-                print(f"validate failed on node {r}.\n Reason: {reason}, {e}")
-            raise
-        if config.scan__debug:
-            from aesara.scan.op import Scan
-
-            scans2 = [n for n in fgraph.apply_nodes if isinstance(n.op, Scan)]
-            nb = len(scans)
-            nb2 = len(scans2)
-            if nb2 > nb:
                 print(
-                    "Extra scan introduced",
-                    nb,
-                    nb2,
-                    getattr(reason, "name", reason),
-                    r,
-                    new_r,
-                )
-            elif nb2 < nb:
-                print(
-                    "Scan removed", nb, nb2, getattr(reason, "name", reason), r, new_r
+                    f"optimizer: validate failed on node {r}.\n Reason: {reason}, {e}"
                 )
+            raise
+
         if verbose:
-            print(reason, r, new_r)
+            print(f"optimizer: rewrite {reason} replaces {r} with {new_r}")
+
         # The return is needed by replace_all_validate_remove
         return chk
 

aesara/graph/fg.py

@@ -510,7 +510,7 @@ def replace(
         if verbose is None:
             verbose = config.optimizer_verbose
         if verbose:
-            print(reason, var, new_var)
+            print(f"optimizer: rewrite {reason} replaces {var} with {new_var}")
 
         new_var = var.type.filter_variable(new_var, allow_convert=True)
 

aesara/graph/opt.py

@@ -1347,6 +1347,10 @@ def transform(self, fgraph, node):
                     new_vars = new_repl
                 else:  # It must be a dict
                     new_vars = list(new_repl.values())
+
+                if config.optimizer_verbose:
+                    print(f"optimizer: rewrite {opt} replaces {node} with {new_repl}")
+
                 if self.profile:
                     self.node_created[opt] += len(
                         list(applys_between(fgraph.variables, new_vars))

aesara/ifelse.py

@@ -634,7 +634,6 @@ def apply(self, fgraph):
                     gpu=False,
                     name=mn_name + "&" + pl_name,
                 )
-                print("here")
                 new_outs = new_ifelse(*new_ins, return_list=True)
                 new_outs = [clone_replace(x) for x in new_outs]
                 old_outs = []

aesara/tensor/basic.py

@@ -8,7 +8,6 @@
 import builtins
 import logging
 import warnings
-from collections import OrderedDict
 from collections.abc import Sequence
 from functools import partial
 from numbers import Number
@@ -26,7 +25,7 @@
 from aesara.graph.basic import Apply, Constant, Variable
 from aesara.graph.op import COp, Op
 from aesara.graph.params_type import ParamsType
-from aesara.graph.type import CType
+from aesara.graph.type import Type
 from aesara.misc.safe_asarray import _asarray
 from aesara.printing import min_informative_str, pprint
 from aesara.scalar import int32
@@ -38,7 +37,7 @@
     get_vector_length,
 )
 from aesara.tensor.elemwise import DimShuffle, Elemwise, scalar_elemwise
-from aesara.tensor.exceptions import EmptyConstantError, NotScalarConstantError
+from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.shape import (
     Shape,
     Shape_i,
@@ -121,7 +120,7 @@ def _as_tensor_Scalar(x, name, ndim, **kwargs):
 def _as_tensor_Variable(x, name, ndim, **kwargs):
     if not isinstance(x.type, TensorType):
         raise TypeError(
-            "Tensor type field must be a TensorType; found {}.".format(type(x.type))
+            f"Tensor type field must be a TensorType; found {type(x.type)}."
         )
 
     if ndim is None:
@@ -135,9 +134,7 @@ def _as_tensor_Variable(x, name, ndim, **kwargs):
         x = x.dimshuffle(list(range(x.ndim))[first_non_broadcastable:])
         if x.ndim > ndim:
             raise ValueError(
-                "Tensor of type {} could not be cast to have {} dimensions".format(
-                    x.type, ndim
-                )
+                f"Tensor of type {x.type} could not be cast to have {ndim} dimensions"
             )
         return x
     elif x.type.ndim < ndim:
@@ -258,31 +255,6 @@ def _obj_is_wrappable_as_tensor(x):
         return False
 
 
-def numpy_scalar(data):
-    """Return a scalar stored in a numpy ndarray.
-
-    Raises
-    ------
-    NotScalarConstantError
-        If the numpy ndarray is not a scalar.
-    EmptyConstantError
-
-    """
-
-    # handle case where data is numpy.array([])
-    if data.ndim > 0 and (len(data.shape) == 0 or builtins.max(data.shape) == 0):
-        assert np.all(np.array([]) == data)
-        raise EmptyConstantError()
-    try:
-        complex(data)  # works for all numeric scalars
-        return data
-    except Exception:
-        raise NotScalarConstantError(
-            "v.data is non-numeric, non-scalar, or has more than one" " unique value",
-            data,
-        )
-
-
 _scalar_constant_value_elemwise_ops = (
     aes.Cast,
     aes.Switch,
@@ -345,7 +317,10 @@ def get_scalar_constant_value(
             return np.asarray(v)
 
         if isinstance(v, np.ndarray):
-            return numpy_scalar(v).copy()
+            try:
+                return np.array(v.item(), dtype=v.dtype)
+            except ValueError:
+                raise NotScalarConstantError()
 
         if isinstance(v, Constant):
             if getattr(v.tag, "unique_value", None) is not None:
@@ -354,7 +329,10 @@ def get_scalar_constant_value(
                 data = v.data
 
             if isinstance(data, np.ndarray):
-                return numpy_scalar(data).copy()
+                try:
+                    return np.array(data.item(), dtype=v.dtype)
+                except ValueError:
+                    raise NotScalarConstantError()
             else:
                 return data
 
@@ -461,28 +439,14 @@ def get_scalar_constant_value(
                     and isinstance(v.owner.inputs[0].owner.op, Join)
                     and len(v.owner.op.idx_list) == 1
                 ):
-                    # Ensure the Join is joining only scalar variables (so that
-                    # the constant value can be found at the same index as the
-                    # one used in the sub-tensor).
-                    if builtins.all(
-                        var.ndim == 0 for var in v.owner.inputs[0].owner.inputs[1:]
-                    ):
-                        idx = v.owner.op.idx_list[0]
-                        if isinstance(idx, CType):
-                            idx = get_scalar_constant_value(
-                                v.owner.inputs[1], max_recur=max_recur
-                            )
-                        # Note the '+ 1' is because the first argument to Join
-                        # is the axis.
-                        ret = v.owner.inputs[0].owner.inputs[idx + 1]
-                        ret = get_scalar_constant_value(ret, max_recur=max_recur)
-                        # join can cast implicitly its input in some case.
-                        return _asarray(ret, dtype=v.type.dtype)
+                    # Ensure the Join is joining only (effectively) scalar
+                    # variables (so that the constant value can be found at the
+                    # same index as the one used in the sub-tensor).
                     if builtins.all(
                         var.ndim == 1 for var in v.owner.inputs[0].owner.inputs[1:]
                     ):
                         idx = v.owner.op.idx_list[0]
-                        if isinstance(idx, CType):
+                        if isinstance(idx, Type):
                             idx = get_scalar_constant_value(
                                 v.owner.inputs[1], max_recur=max_recur
                             )
@@ -517,7 +481,7 @@ def get_scalar_constant_value(
                 ):
 
                     idx = v.owner.op.idx_list[0]
-                    if isinstance(idx, CType):
+                    if isinstance(idx, Type):
                         idx = get_scalar_constant_value(
                             v.owner.inputs[1], max_recur=max_recur
                         )
@@ -539,7 +503,7 @@ def get_scalar_constant_value(
                     op = owner.op
                     idx_list = op.idx_list
                     idx = idx_list[0]
-                    if isinstance(idx, CType):
+                    if isinstance(idx, Type):
                         idx = get_scalar_constant_value(
                             owner.inputs[1], max_recur=max_recur
                         )
@@ -576,12 +540,7 @@ def get_scalar_constant_value(
                     if isinstance(grandparent, Constant):
                         return np.asarray(np.shape(grandparent.data)[idx])
 
-        raise NotScalarConstantError(v)
-
-
-#########################
-# Casting Operations
-#########################
+        raise NotScalarConstantError()
 
 
 class TensorFromScalar(Op):
@@ -702,7 +661,7 @@ class Rebroadcast(COp):
     def __init__(self, *axis):
         # Sort them to make sure we merge all possible case.
         items = sorted(axis)
-        self.axis = OrderedDict(items)
+        self.axis = dict(items)
         for axis, broad in self.axis.items():
             if not isinstance(axis, (np.integer, int)):
                 raise TypeError(f"Rebroadcast needs integer axes. Got {axis}")
@@ -719,13 +678,7 @@ def __hash__(self):
         return hash((type(self), tuple(items)))
 
     def __str__(self):
-        if len(self.axis) == 0:
-            broadcast_pattern = []
-        else:
-            broadcast_pattern = ["?" for i in range(1 + max(self.axis.keys()))]
-        for k, v in self.axis.items():
-            broadcast_pattern[k] = str(int(v))
-        return f"{self.__class__.__name__}{{{','.join(broadcast_pattern)}}}"
+        return f"{self.__class__.__name__}{{{','.join(str(i) for i in self.axis.items())}}}"
 
     def make_node(self, x):
         if self.axis.keys() and (x.ndim <= max(self.axis.keys())):