small fixes

Author: dougalm

jax/_src/interpreters/partial_eval.py

@@ -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 = {}
@@ -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):