Fix reference cycle

examples/neural_y.py

@@ -13,17 +13,15 @@ def __init__(self, f, approx_fix_f):
         self.approx_fix_f = approx_fix_f
         self.buffers.loss = 0  # FIXME
 
-        def approx_f(_, *args):
-            return self.approx_fix_f(*args)
-
-        self.approx_f = approx_f
+    def approx_f(self, fn, *args):
+        return self.approx_fix_f(*args)
 
     def forward(self, *args):
+        # TODO Unroll to multiple depths, gathering loss at each depth.
         value0 = self.approx_fix_f(None, *args)
         value1 = self.f(self.approx_f, *args)
         self.buffers.loss = (value0 - value1).abs().sum()
-        value = 0.5 * (value0 + value1)
-        return value
+        return value1
 
 
 def main(args):

funsor/fixpoints.py

@@ -20,12 +20,12 @@
 def truncate(fn, output):
     TRUNCATION_DEPTHS[fn] += 1
     if TRUNCATION_DEPTHS[fn] <= MAX_TRUNCATION_DEPTH:
-        result = fn
+        approx_fn = fn
     else:
-        def result(*args, **kwargs):
+        def approx_fn(fn, *args):
             return Tensor(torch.tensor(float('nan')).expand(output.shape))
     try:
-        yield result
+        yield approx_fn
     finally:
         TRUNCATION_DEPTHS[fn] -= 1
 
@@ -74,8 +74,8 @@ def eager_subs(self, subs):
 
 @truncated.register(Fix, object, Domain, tuple)
 def eager_function(fn, output, args):
-    with truncate(fn) as t:
-        return fn(t, *args)
+    with truncate(fn) as approx_fn:
+        return fn(approx_fn, *args)
 
 
 def _fix(inputs, output, fn):