Make TF meta objects follow their base APIs more closely

* More memoization/caching has been added along with a small amount of
refactoring to the meta API.
* Hashing for special cases (e.g. Numpy arrays and dicts) has been moved closer
to their sources (e.g. constant tensor classes and TF NodeDef).
* Superfluous TF Python meta types have been removed.
* The name property for TF Tensors is now derived directly from the name field
in the corresponding NodeDef.
* Dispatch functions now have distinct names (good for debugging).
* metatize for types/classes returns a meta type/class, not an
instance of the type/class

Author: brandonwillard

requirements-dev.txt

@@ -4,3 +4,4 @@ pytest-cov>=2.6.1
 pytest-html>=1.20.0
 pylint>=2.3.1
 black>=19.3b0
+ipython

requirements.txt

@@ -11,3 +11,4 @@ unification>=0.2.2
 kanren @ git+https://github.com/pymc-devs/kanren.git@symbolic-pymc#egg=kanren-0.2.3
 toolz>=0.9.0
 sympy>=1.3
+cachetools

symbolic_pymc/meta.py

@@ -2,8 +2,6 @@
 import types
 import reprlib
 
-import numpy as np
-
 from itertools import chain
 from functools import partial
 from collections import OrderedDict
@@ -15,11 +13,15 @@
 
 from multipledispatch import dispatch
 
+from cachetools import cached
+
 meta_repr = reprlib.Repr()
 meta_repr.maxstring = 100
 meta_repr.maxother = 100
 meta_repr.print_obj = False
 
+metatize_cache = {}
+
 
 def metatize(obj):
     """Convert object to base type then meta object."""
@@ -28,38 +30,39 @@ def metatize(obj):
     return _metatize(obj)
 
 
-@dispatch((set, list, tuple))
+@dispatch((type(None), types.FunctionType, partial, str, dict))
 def _metatize(obj):
+    return obj
+
+
+@_metatize.register((set, tuple))
+@cached(metatize_cache)
+def _metatize_set_tuple(obj):
     """Convert elements of an iterable to meta objects."""
     return type(obj)([metatize(o) for o in obj])
 
 
-@dispatch(Iterator)
-def _metatize(obj):
-    """Convert elements of an iterator to meta objects."""
-    return iter([metatize(o) for o in obj])
+@_metatize.register(list)
+def _metatize_list(obj):
+    """Convert elements of an iterable to meta objects."""
+    return type(obj)([metatize(o) for o in obj])
 
 
-def _make_hashable(x):
-    if isinstance(x, list):
-        return tuple(x)
-    elif isinstance(x, Mapping):
-        return frozenset(x.items())
-    elif isinstance(x, np.ndarray):
-        return x.tostring()
-    else:
-        return x
+@_metatize.register(Iterator)
+@cached(metatize_cache)
+def _metatize_Iterator(obj):
+    """Convert elements of an iterator to meta objects."""
+    return iter([metatize(o) for o in obj])
 
 
-def _meta_reify_iter(rands):
+def meta_reify_iter(rands):
     """Recursively reify an iterable object and return a boolean indicating the presence of un-reifiable objects, if any."""
-    # We want as many of the rands reified as possible,
     any_unreified = False
     reified_rands = []
-    if isinstance(rands, Mapping):
-        _rands = rands.items()
-    else:
-        _rands = rands
+
+    _rands = rands
+    if isinstance(_rands, Mapping):
+        _rands = _rands.items()
 
     for s in _rands:
         if isinstance(s, MetaSymbol):
@@ -71,11 +74,11 @@ def _meta_reify_iter(rands):
             reified_rands.append(s)
             any_unreified |= True
         elif isinstance(s, (list, tuple)):
-            _reified_rands, _any_unreified = _meta_reify_iter(s)
+            _reified_rands, _any_unreified = meta_reify_iter(s)
             reified_rands.append(type(s)(_reified_rands))
             any_unreified |= _any_unreified
         else:
-            reified_rands += [s]
+            reified_rands.append(s)
 
     return type(rands)(reified_rands), any_unreified
 
@@ -153,8 +156,6 @@ def _cached_hash(self):
 
             new_cls.__hash__ = _cached_hash
 
-        # TODO: Could register base classes.
-        # E.g. cls.register(bases)
         return new_cls
 
 
@@ -164,7 +165,7 @@ class MetaSymbol(metaclass=MetaSymbolType):
     TODO: Should `MetaSymbol.obj` be an abstract property and a `weakref`?
     """
 
-    __slots__ = ("_obj", "_hash")
+    __slots__ = ("_obj", "_hash", "_rands")
 
     @property
     @abc.abstractmethod
@@ -189,18 +190,40 @@ def is_meta(cls, obj):
         return isinstance(obj, MetaSymbol) or isvar(obj)
 
     def __init__(self, obj=None):
+        assert obj is None or isvar(obj) or isinstance(obj, self.base)
         self._obj = obj
 
     def rands(self):
-        """Create a tuple of the meta object's operator parameters (i.e. "rands")."""
-        return tuple(getattr(self, s) for s in getattr(self, "__all_props__", ()))
+        """Get a tuple of the meta object's operator parameters (i.e. "rands")."""
+        if getattr(self, "_rands", None) is not None:
+            return self._rands
+
+        self._rands = tuple(getattr(self, s) for s in getattr(self, "__all_props__", ()))
+
+        return self._rands
 
     def reify(self):
-        """Create a concrete base object from this meta object (and its rands)."""
+        """Attempt to create a concrete base object from this meta object.
+
+        During the process, dependent objects will need to be reified, which
+        may result in updates to the object(s) being reified.
+
+        For instance, if a meta tensor's parent operator is fully reifiable to
+        a base object, then the meta tensor's dtype and shape may be fixed:
+        e.g. a tensor corresponding to the output of a sum of two float64
+        scalars is necessarily a float64 scalar.
+
+        This function will set any unspecified properties (e.g. dtype and shape
+        values for the previous example), mutating the object in-place when
+        possible.  It will return a [refined/partially reified] meta object
+        when it can't fully reify to a base object (in which case, it will
+        return the base object) or when partial reification results in a meta
+        object from a subclass.
+        """
         if self.obj is not None and not isinstance(self.obj, Var):
             return self.obj
         else:
-            reified_rands, any_unreified = _meta_reify_iter(self.rands())
+            reified_rands, any_unreified = meta_reify_iter(self.rands())
 
             # If not all the rands reified, then create another meta
             # object--albeit one with potentially more non-`None` `obj` fields.
@@ -220,35 +243,20 @@ def __eq__(self, other):
         if not (type(self) == type(other)):
             return False
 
-        if not (self.base == other.base):
-            return False
+        assert self.base == other.base
 
-        a_slots = getattr(self, "__all_props__", None)
-        if a_slots is not None:
-            if not all(_check_eq(getattr(self, attr), getattr(other, attr)) for attr in a_slots):
-                return False
-        elif getattr(other, "__all_props__", None) is not None:
-            # The other object has slots, but this one doesn't.
-            return False
+        if self.rands():
+            return all(_check_eq(s, o) for s, o in zip(self.rands(), other.rands()))
         else:
-            # Neither have slots, so best we can do is compare
-            # base objects (if any).
-            # If there aren't base objects, we say they're not equal.
-            # (Maybe we should *require* base objects in this case
-            # and raise an exception?)
-            return getattr(self, "obj", None) == getattr(other, "obj", None) is not None
+            return NotImplemented
 
-        return True
+        return False
 
     def __ne__(self, other):
         return not self.__eq__(other)
 
     def __hash__(self):
-        if getattr(self, "__props__", None) is not None:
-            rands = tuple(_make_hashable(p) for p in self.rands())
-            return hash(rands + (self.base,))
-        else:
-            return hash((self.base, self.obj))
+        return hash((self.base, self.rands()))
 
     def __str__(self):
         obj = getattr(self, "obj", None)
@@ -273,8 +281,8 @@ def _repr_pretty_(self, p, cycle):
             with p.group(2, f"{self.__class__.__name__}(", ")"):
                 p.breakable()
                 idx = None
-                if hasattr(self, "__props__"):
-                    for idx, (name, item) in enumerate(zip(self.__props__, self.rands())):
+                if hasattr(self, "__all_props__"):
+                    for idx, (name, item) in enumerate(zip(self.__all_props__, self.rands())):
                         if idx:
                             p.text(",")
                             p.breakable()
@@ -292,8 +300,8 @@ def _repr_pretty_(self, p, cycle):
                     p.pretty(obj)
 
 
-@dispatch((MetaSymbol, type(None), types.FunctionType, partial, str, dict))
-def _metatize(obj):
+@_metatize.register(MetaSymbol)
+def _metatize_MetaSymbol(obj):
     return obj
 
 
@@ -314,21 +322,26 @@ class MetaOp(MetaSymbol):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    @MetaSymbol.obj.setter
-    def obj(self, x):
-        if hasattr(self, "_obj"):
-            raise ValueError("Cannot reset obj in an `Op`")
-        object.__setattr__(self, "_obj", x)
-
     @abc.abstractmethod
     def out_meta_types(self, inputs=None):
         """Return the types of meta variables this `Op` is expected to produce given the inputs."""
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def __call__(self, *args, ttype=None, index=None, **kwargs):
+    def __call__(self, *args, **kwargs):
         raise NotImplementedError()
 
+    def __eq__(self, other):
+        res = super().__eq__(other)
+
+        if res is NotImplemented:
+            return getattr(self, "obj", None) == getattr(other, "obj", None) is not None
+
+        return res
+
+    def __hash__(self):
+        return hash((self.base, self.obj))
+
 
 class MetaVariable(MetaSymbol):
     __slots__ = ()
@@ -369,14 +382,25 @@ def _find_meta_type(obj_type, meta_abs_type):
             # This object is a subclass of an existing meta class' base type,
             # but there is no implemented meta type for this subclass, so we
             # dynamically make one.
+
+            # FIXME, TODO: We should do something about `Op` constructor
+            # arguments and properties.
+            #
+            # For instance, `tt.nlinalg.SVD` takes `full_matrices` and `compute_uv`
+            # constructor arguments, but the dynamically constructed `TheanoMetaOp` type for
+            # SVD is just the base `TheanoMetaOp.__init__`, which doesn't account for those.
+            # To do this correctly, we would need to dynamically metatize the underlying
+            # `Op`'s `__init__` and so on.
             new_type = type(f"Meta{obj_type.__name__}", (obj_cls,), {"base": obj_type})
-            return new_type(obj_type)
+
+            return new_type
         else:
             cls = obj_cls
 
 
-@dispatch(type)
-def _metatize(obj_type):
+@_metatize.register(type)
+@cached(metatize_cache)
+def _metatize_type(obj_type):
     """Return an existing meta type/class, or create a new one."""
     for meta_type in MetaSymbol.__subclasses__():
         obj_cls = _find_meta_type(obj_type, meta_type)