Merge pull request #69 from brandonwillard/really-fix-numpy-unify

Fix unify when walked values are Numpy arrays

Author: brandonwillard

requirements-dev.txt

@@ -1,5 +1,5 @@
 pydocstyle>=3.0.0
-pytest>=4.2.0
+pytest>=5.0.0
 pytest-cov>=2.6.1
 pytest-html>=1.20.0
 pylint>=2.3.1

symbolic_pymc/tensorflow/meta.py

@@ -48,16 +48,14 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @classmethod
-    def make_opdef_sig(cls, opdef):
+    def make_opdef_sig(cls, opdef, opdef_py_func=None):
         """Create a `Signature` object for an `OpDef`.
 
         Annotations are include so that one can partially verify arguments.
         """
         input_args = OrderedDict([(a.name, a.type or a.type_attr) for a in opdef.input_arg])
         attrs = OrderedDict([(a.name, a.type) for a in opdef.attr])
 
-        opdef_py_func = getattr(tf.raw_ops, opdef.name, None)
-
         params = OrderedDict()
         if opdef_py_func:
             # We assume we're dealing with a function from `tf.raw_ops`.
@@ -86,6 +84,9 @@ def make_opdef_sig(cls, opdef):
                 params[name] = new_param
 
         else:
+            # We're crafting the Operation from a low-level via `apply_op`.
+            opdef_py_func = partial(op_def_lib.apply_op, opdef.name)
+
             for i_name, i_type in input_args.items():
                 p = Parameter(i_name, Parameter.POSITIONAL_OR_KEYWORD, annotation=i_type)
                 params[i_name] = p
@@ -117,15 +118,17 @@ def make_opdef_sig(cls, opdef):
         opdef_sig = Signature(
             params.values(), return_annotation=[(o.name, o.type_attr) for o in opdef.output_arg]
         )
-        return opdef_sig
+        return opdef_sig, opdef_py_func
 
     def add_op(self, opdef):
         op_info = self._ops.get(opdef.name, None)
         if op_info is None:
             super().add_op(opdef)
             op_info = self._ops[opdef.name]
-            opdef_sig = self.make_opdef_sig(op_info.op_def)
+            opdef_func = getattr(tf.raw_ops, opdef.name, None)
+            opdef_sig, opdef_func = self.make_opdef_sig(op_info.op_def, opdef_func)
             op_info.opdef_sig = opdef_sig
+            op_info.opdef_func = opdef_func
         return op_info
 
     def get_opinfo(self, opdef):
@@ -239,7 +242,7 @@ class TFlowMetaOpDef(MetaOp, TFlowMetaSymbol):
     def __init__(self, obj=None):
         op_info = op_def_lib.add_op(obj)
         self.apply_func_sig = op_info.opdef_sig
-        self.apply_func = partial(op_def_lib.apply_op, obj.name)
+        self.apply_func = op_info.opdef_func
         super().__init__(obj=obj)
 
     def out_meta_types(self, inputs=None):

symbolic_pymc/unify.py

@@ -43,23 +43,31 @@ def wrapper(*args, **kwargs):
         _unify._cache.clear()
 
 
+_base_unify = unify.dispatch(object, object, dict)
+
+
 def unify_numpy(u, v, s):
     """Handle NumPy arrays in a special way to avoid warnings/exceptions."""
+    u = walk(u, s)
     v = walk(v, s)
+
+    if u is v:
+        return s
     if isvar(u):
         return assoc(s, u, v)
     if isvar(v):
         return assoc(s, v, u)
-    # Switch the order of comparison so that `v.__eq__` is tried (in case it's
-    # not also a NumPy array, but has logic for such comparisons)
-    if np.array_equal(v, u):
+
+    if isinstance(u, np.ndarray) or isinstance(v, np.ndarray):
+        if np.array_equal(v, u):
+            return s
+    elif u == v:
         return s
+
     return _unify(u, v, s)
 
 
-unify.add((np.ndarray, object, dict), unify_numpy)
-unify.add((object, np.ndarray, dict), unify_numpy)
-unify.add((np.ndarray, np.ndarray, dict), unify_numpy)
+unify.add((object, object, dict), unify_numpy)
 
 
 def unify_MetaSymbol(u, v, s):

tests/tensorflow/test_meta.py

@@ -167,6 +167,7 @@ def test_meta_create():
     with pytest.raises(TypeError):
         TFlowMetaTensor('float64', 'Add', name='q__')
 
+
 @pytest.mark.usefixtures("run_with_tensorflow")
 @run_in_graph_mode
 def test_meta_Op():
@@ -196,7 +197,6 @@ def test_meta_Op():
         assert MetaSymbol.is_meta(test_op.outputs[0])
 
 
-
 @pytest.mark.usefixtures("run_with_tensorflow")
 def test_meta_lvars():
     """Make sure we can use lvars as values."""
@@ -381,7 +381,7 @@ def test_opdef_sig():
 
     custom_opdef_tf.attr.extend([attr1_tf, attr2_tf])
 
-    opdef_sig = MetaOpDefLibrary.make_opdef_sig(custom_opdef_tf)
+    opdef_sig, opdef_func = MetaOpDefLibrary.make_opdef_sig(custom_opdef_tf)
 
     import inspect
     # These are standard inputs
@@ -431,3 +431,13 @@ class CustomClass(object):
 
     with pytest.raises(ValueError):
         mt(CustomClass())
+
+
+@pytest.mark.usefixtures("run_with_tensorflow")
+@run_in_graph_mode
+def test_opdef_func():
+    sum_mt = mt.Sum([[1, 2]], [1])
+    sum_tf = sum_mt.reify()
+
+    with tf.compat.v1.Session() as sess:
+        assert sum_tf.eval() == np.r_[3]

tests/test_unify.py

@@ -36,3 +36,15 @@ def test_numpy():
     s = unify([1, var('a')], np_array)
 
     assert s is False
+
+    s = unify(var('a'), 2, {var('a'): np_array})
+
+    assert s is False
+
+    s = unify(var('a'), var('b'), {var('a'): np_array})
+
+    assert s[var('a')] is np_array
+    assert s[var('b')] is np_array
+
+    s = unify(np_array, np_array)
+    assert s == {}