Ensure alpha conversion always runs when a new variable is bound

funsor/terms.py

@@ -41,19 +41,12 @@ def subs_interpreter(cls, *args):
         return interpreter.reinterpret(expr)
 
 
-def _alpha_mangle(expr):
+def _alpha_mangle(bound_vars):
     """
     Rename bound variables in expr to avoid conflict with any free variables.
-
-    FIXME this does not avoid conflict with other bound variables.
+    Returns substitution dictionary with mangled names for consumption by Funsor._alpha_convert.
     """
-    alpha_subs = {name: interpreter.gensym(name + "__BOUND")
-                  for name in expr.bound if "__BOUND" not in name}
-    if not alpha_subs:
-        return expr
-
-    ast_values = expr._alpha_convert(alpha_subs)
-    return reflect(type(expr), *ast_values)
+    return {name: interpreter.gensym(name.split("__BOUND_")[0] + "__BOUND_") for name in bound_vars}
 
 
 def reflect(cls, *args, **kwargs):
@@ -81,8 +74,24 @@ def reflect(cls, *args, **kwargs):
     result = super(FunsorMeta, cls_specific).__call__(*args)
     result._ast_values = args
 
-    # alpha-convert eagerly upon binding any variable
-    result = _alpha_mangle(result)
+    # alpha-convert eagerly upon binding any variable.
+    # the identifier we use to reconcile alpha-conversion and cons-hashing
+    # is the string literal of hash() of the type and cons-hashing key:
+    if result.bound:
+        alpha_subs = _alpha_mangle(result.bound)
+        alpha_mangled_args = result._alpha_convert(alpha_subs)
+
+        # TODO eliminate code duplication below
+        # this is currently necessary because .bound is computed in __init__().
+        result = super(FunsorMeta, cls_specific).__call__(*alpha_mangled_args)
+        result._ast_values = alpha_mangled_args
+
+        # we also make the old cons cache_key point to the new mangled value.
+        # this guarantees that alpha-conversion only runs once for this expression.
+        cls._cons_cache[cache_key] = result
+
+        cache_key = tuple(id(arg) if type(arg).__name__ == "DeviceArray" or not isinstance(arg, Hashable)
+                          else arg for arg in alpha_mangled_args)
 
     cls._cons_cache[cache_key] = result
     return result

test/test_alpha_conversion.py

@@ -115,3 +115,11 @@ def test_sample_independent():
     actual = Independent(f, 'x', 'i', 'x_i')
     assert actual.sample('i')
     assert actual.sample('j', {'i': 2})
+
+
+def test_subs_already_bound():
+    with interpretation(reflect):
+        x = Variable('x', Real)
+        y1 = (2 * x)(x=3)
+        y2 = y1.arg(4)
+    assert y1.bound != y2.bound