Merge pull request #85 from brandonwillard/custom-ndarray-class

Simplify NumPy array handling by using a custom ndarray class

Author: brandonwillard

conftest.py

@@ -1,18 +1,16 @@
 import pytest
 
+
 @pytest.fixture()
 def run_with_theano():
-    from symbolic_pymc.theano.meta import load_dispatcher
+    # import symbolic_pymc.meta
 
-    load_dispatcher()
+    # from symbolic_pymc.meta import base_metatize
 
+    import symbolic_pymc.theano.meta as tm
 
-@pytest.fixture()
-def run_with_tensorflow():
-    from symbolic_pymc.tensorflow.meta import load_dispatcher
+    tm.load_dispatcher()
 
-    load_dispatcher()
+    # yield
 
-    # Let's make sure we have a clean graph slate
-    from tensorflow.compat.v1 import reset_default_graph
-    reset_default_graph()
+    # symbolic_pymc.meta._metatize = base_metatize

symbolic_pymc/meta.py

@@ -2,6 +2,9 @@
 import types
 import reprlib
 
+import numpy as np
+
+from copy import deepcopy
 from itertools import chain
 from functools import partial
 from collections import OrderedDict
@@ -10,7 +13,7 @@
 
 from unification import isvar, Var
 
-from .utils import _check_eq
+from .utils import HashableNDArray
 
 from multipledispatch import dispatch
 
@@ -89,6 +92,18 @@ def _metatize_list(obj):
     return type(obj)([metatize(o) for o in obj])
 
 
+@_metatize.register(np.ndarray)
+def _metatize_ndarray(obj):
+    """Convert Numpy ndarrays into hashable objects."""
+    return metatize(obj.view(HashableNDArray))
+
+
+@_metatize.register(HashableNDArray)
+def _metatize_HashableNDArray(obj):
+    """Fallback case for converted Numpy ndarrays."""
+    return obj
+
+
 @_metatize.register(Iterator)
 @cached(metatize_cache)
 def _metatize_Iterator(obj):
@@ -183,9 +198,6 @@ def __metatize(cls, obj):
 
         new_cls = super().__new__(cls, name, bases, clsdict)
 
-        if isinstance(new_cls.base, type):
-            _metatize.add((new_cls.base,), new_cls._metatize)
-
         # Wrap the class implementation of `__hash__` with this value-caching
         # code.
         if "_hash" in clsdict["__volatile_slots__"]:
@@ -226,12 +238,11 @@ def obj(self):
         return object.__getattribute__(self, "_obj")
 
     @classmethod
-    def base_classes(cls, mro_order=True):
-        res = tuple(c.base for c in cls.__subclasses__())
-        if hasattr(cls, "base"):
-            res = (cls.base,) + res
-        sorted(res, key=lambda c: len(c.mro()), reverse=mro_order)
-        return res
+    def base_subclasses(cls):
+        for subclass in cls.__subclasses__():
+            yield from subclass.base_subclasses()
+            if isinstance(subclass.base, type):
+                yield subclass
 
     @classmethod
     def is_meta(cls, obj):
@@ -294,7 +305,7 @@ def __eq__(self, other):
         assert self.base == other.base
 
         if self.rands():
-            return all(_check_eq(s, o) for s, o in zip(self.rands(), other.rands()))
+            return all(s == o for s, o in zip(self.rands(), other.rands()))
         else:
             return NotImplemented
 
@@ -459,3 +470,6 @@ def _metatize_type(obj_type):
 
         if obj_cls is not None:
             return obj_cls
+
+
+base_metatize = deepcopy(_metatize)

symbolic_pymc/relations/theano/distributions.py

@@ -10,6 +10,7 @@
 
 from .. import concat
 from ...etuple import etuple
+from ...utils import HashableNDArray
 from ...theano.meta import TheanoMetaConstant, mt
 
 from kanren.facts import Relation
@@ -61,18 +62,18 @@ def constant_neq(lvar, val):
     Scalar values are broadcast across arrays.
     """
 
-    def _goal(s):
+    if isinstance(val, np.ndarray):
+        val = val.view(HashableNDArray)
+
+    def constant_neq_goal(s):
         lvar_val = walk(lvar, s)
         if isinstance(lvar_val, (tt.Constant, TheanoMetaConstant)):
-            data = lvar_val.data
-            if (isinstance(val, np.ndarray) and not np.array_equal(data, val)) or not all(
-                np.atleast_1d(data) == val
-            ):
+            if lvar_val.data != val:
                 yield s
         else:
             yield s
 
-    return _goal
+    return constant_neq_goal
 
 
 def scale_loc_transform(in_expr, out_expr):

symbolic_pymc/tensorflow/meta.py

@@ -1,8 +1,6 @@
 import types
 import inspect
 
-import numpy as np
-
 import tensorflow as tf
 import tensorflow_probability as tfp
 
@@ -39,6 +37,9 @@
 
 from .. import meta
 
+from ..utils import HashableNDArray
+
+
 tf_metatize_cache = Cache(50)
 
 
@@ -167,7 +168,7 @@ def _metatize_tf_object(obj):
     try:
         tf_obj = tf.convert_to_tensor(obj)
     except (TypeError, ValueError):
-        raise ValueError("Could not find a TensorFlow MetaSymbol class for {obj}")
+        raise ValueError(f"Error converting {obj} to a TensorFlow tensor.")
 
     return _metatize(tf_obj)
 
@@ -180,9 +181,9 @@ def load_dispatcher():
 
     def _metatize_tf_svd(obj):
         """Turn a TensorFlow `Svd` object/tuple into a standard tuple."""
-        return _metatize(tuple(obj))
+        return meta._metatize(tuple(obj))
 
-    _metatize.add((_SvdOutput,), _metatize_tf_svd)
+    meta._metatize.add((_SvdOutput,), _metatize_tf_svd)
 
     def _metatize_tf_eager(obj):
         """Catch eager tensor metatize issues early."""
@@ -192,12 +193,17 @@ def _metatize_tf_eager(obj):
             " (e.g. within `tensorflow.python.eager.context.graph_mode`)"
         )
 
-    _metatize.add((EagerTensor,), _metatize_tf_eager)
+    meta._metatize.add((EagerTensor,), _metatize_tf_eager)
 
-    _metatize.add((object,), _metatize_tf_object)
+    meta._metatize.add((object,), _metatize_tf_object)
+    meta._metatize.add((HashableNDArray,), _metatize_tf_object)
 
+    for new_cls in TFlowMetaSymbol.base_subclasses():
+        meta._metatize.add((new_cls.base,), new_cls._metatize)
 
-load_dispatcher()
+    meta._metatize.add((TFlowMetaOpDef.base,), TFlowMetaOpDef._metatize)
+
+    return meta._metatize
 
 
 class TFlowMetaSymbol(MetaSymbol):
@@ -451,7 +457,7 @@ def _protobuf_convert(cls, k, v):
         elif k == "T":
             return tf.as_dtype(v.type).name
         elif k == "value":
-            return tensor_util.MakeNdarray(v.tensor)
+            return tensor_util.MakeNdarray(v.tensor).view(HashableNDArray)
         else:
             # Consider only the narrow case where a single object is converted
             # (e.g. a Python builtin type under `v.b`, `v.f`, etc.)
@@ -492,9 +498,7 @@ def frozen_attr(self):
         if isvar(self.attr):
             self._frozen_attr = self.attr
         else:
-            self._frozen_attr = frozenset(
-                (k, v.tostring() if isinstance(v, np.ndarray) else v) for k, v in self.attr.items()
-            )
+            self._frozen_attr = frozenset(self.attr.items())
         return self._frozen_attr
 
     def __eq__(self, other):
@@ -993,3 +997,5 @@ def __getattr__(self, obj):
 
 
 mt = TFlowMetaAccessor()
+
+load_dispatcher()

symbolic_pymc/theano/meta.py

@@ -23,21 +23,35 @@
     _metatize,
 )
 
+from .. import meta
+
+from ..utils import HashableNDArray
+
 
 def _metatize_theano_object(obj):
     try:
         obj = tt.as_tensor_variable(obj)
     except (ValueError, tt.AsTensorError):
-        raise ValueError("Could not find a MetaSymbol class for {}".format(obj))
+        raise ValueError("Error converting {} to a Theano tensor.".format(obj))
+    except AssertionError:
+        # This is a work-around for a Theano bug; specifically,
+        # an assert statement in `theano.scalar.basic` that unnecessarily
+        # requires the object type be exclusively an ndarray or memmap.
+        # See https://github.com/Theano/Theano/pull/6727
+        obj = tt.as_tensor_variable(np.asarray(obj))
+
     return _metatize(obj)
 
 
 def load_dispatcher():
     """Set/override dispatcher to default to TF objects."""
-    _metatize.add((object,), _metatize_theano_object)
+    meta._metatize.add((object,), _metatize_theano_object)
+    meta._metatize.add((HashableNDArray,), _metatize_theano_object)
 
+    for new_cls in TheanoMetaSymbol.base_subclasses():
+        meta._metatize.add((new_cls.base,), new_cls._metatize)
 
-load_dispatcher()
+    return meta._metatize
 
 
 class TheanoMetaSymbol(MetaSymbol):
@@ -202,7 +216,7 @@ def __call__(self, *args, ttype=None, index=None, **kwargs):
             # XXX: We don't have a higher-order meta object model, so being
             # wrong about the exact type of output variable will cause
             # problems.
-            out_meta_type, = self.out_meta_types(op_args)
+            (out_meta_type,) = self.out_meta_types(op_args)
             res_var = out_meta_type(ttype, res_apply, index, name)
             res_var._obj = var()
 
@@ -451,28 +465,9 @@ def _metatize(cls, obj):
         return res
 
     def __init__(self, type, data, name=None, obj=None):
-        self.data = data
+        self.data = data if not isinstance(data, np.ndarray) else data.view(HashableNDArray)
         super().__init__(type, None, None, name, obj=obj)
 
-    def __eq__(self, other):
-        if self is other:
-            return True
-
-        if type(self) != type(other):
-            return False
-
-        if all(
-            (s.tostring() if isinstance(s, np.ndarray) else s)
-            == (o.tostring() if isinstance(o, np.ndarray) else o)
-            for s, o in zip(self.rands(), other.rands())
-        ):
-            return True
-
-        return False
-
-    def __hash__(self):
-        return hash(v.tostring() if isinstance(v, np.ndarray) else v for v in self.rands())
-
 
 class TheanoMetaTensorConstant(TheanoMetaConstant):
     # TODO: Could extend `theano.tensor.var._tensor_py_operators`, too.
@@ -606,7 +601,9 @@ def meta_obj(*args, **kwargs):
 
 mt = TheanoMetaAccessor()
 
-mt.dot = metatize(tt.basic._dot)
+_metatize = load_dispatcher()
+
+mt.dot = _metatize(tt.basic._dot)
 
 
 #

symbolic_pymc/theano/ops.py

@@ -168,7 +168,7 @@ def _infer_shape(self, size, dist_params, param_shapes=None):
 
         # _, out_bcasts, bcastd_inputs = tt.add.get_output_info(tt.DimShuffle, *dist_params)
 
-        bcast_ind, = out_bcasts
+        (bcast_ind,) = out_bcasts
         ndim_ind = len(bcast_ind)
         shape_ind = bcastd_inputs[0].shape
 
@@ -226,7 +226,7 @@ def compute_bcast(self, dist_params, size):
                     s_x, s_idx = s.owner.inputs
                     s_idx = tt.get_scalar_constant_value(s_idx)
                     if isinstance(s_x.owner.op, tt.Shape):
-                        x_obj, = s_x.owner.inputs
+                        (x_obj,) = s_x.owner.inputs
                         s_val = x_obj.type.broadcastable[s_idx]
                     else:
                         # TODO: Could go for an existing broadcastable here,

symbolic_pymc/theano/utils.py

@@ -122,7 +122,7 @@ def optimize_graph(x, optimization, return_graph=None, in_place=False):
     else:
         res = x_graph_opt.outputs
         if len(res) == 1:
-            res, = res
+            (res,) = res
     return res