small fixes

Author: dougalm

jax/_src/interpreters/partial_eval.py

@@ -1640,8 +1640,8 @@ class DynamicJaxprTracer(core.Tracer):
   def __init__(self, trace: DynamicJaxprTrace,
                aval: core.AbstractValue | core.AvalQDD,
                val : Atom,
-               parent : TracingEqn | None = None,
-               line_info: source_info_util.SourceInfo | None = None):
+               line_info: source_info_util.SourceInfo | None = None,
+               parent : TracingEqn | None = None):
     # TODO(dougalm): Remove aval. It's redundant now that we have val.
     if isinstance(aval, core.AvalQDD):
       assert aval.qdd is not None
@@ -1767,7 +1767,7 @@ class JaxprStackFrame:
   gensym: Callable[[AbstractValue], Var]
   constid_to_tracer: WeakValueDictionary[ConstId, DynamicJaxprTracer]
   constvar_to_val: dict[Var, Any]
-  tracing_eqns: list[ReferenceType(TracingEqn)]
+  tracing_eqns: list[ReferenceType[TracingEqn]]
   invars: list[Var]
   effects: core.Effects
   debug_info: core.DebugInfo
@@ -1944,6 +1944,7 @@ def __init__(self, debug_info: core.DebugInfo, parent_trace=None, lower=False,
     self.frame = JaxprStackFrame(debug_info, auto_dce)
     self.parent_trace = parent_trace
 
+  # TODO(dougalm): we might be able to remove this since the refcounting should be doing it for us
   def invalidate(self):
     # avoid cyclic refs
     self.frame.constid_to_tracer = {}
@@ -1964,7 +1965,7 @@ def var_to_tracer(self, var, source_info, parent=None):
     aval = var.aval
     if aval.has_qdd:
       aval = core.AvalQDD(aval, var.initial_qdd)
-    return DynamicJaxprTracer(self, aval, var, parent, source_info)
+    return DynamicJaxprTracer(self, aval, var, source_info, parent)
 
   def new_arg(self, aval, source_info: SourceInfo):
     var = self.frame.newvar(aval)
@@ -2022,7 +2023,7 @@ def _new_const(self, aval, c, source_info: SourceInfo) -> DynamicJaxprTracer:
       return tracer
 
   def finalize_const(self, var, constid):
-    self.frame.constvar_to_val.pop(var)
+    self.frame.constvar_to_val.pop(var, None)
 
   def get_const(self, tracer) -> Any:
     atom = tracer.val

tests/api_test.py

@@ -5106,6 +5106,7 @@ def g():
 
 
   def test_implicit_dce(self):
+    @api.jit
     def foo(x):
       const = np.zeros((300,))
       r = weakref.ref(const)
@@ -5114,7 +5115,7 @@ def foo(x):
       assert r() is None, "oops, the constant wasn't DCE'd"
       return x + x
 
-    jax.make_jaxpr(foo, dce=True)(1.0)
+    foo(1.0)
 
 class RematTest(jtu.JaxTestCase):