diff --git a/jax/core.py b/jax/core.py
index <HASH>..<HASH> 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -618,9 +618,6 @@ class MainTrace:
 
 class TraceStack:
   # See comments in https://github.com/google/jax/pull/3370
-  upward: List[MainTrace]
-  downward: List[MainTrace]
-
   def __init__(self):
     eval_trace = MainTrace(0, EvalTrace)
     self.stack = [eval_trace]
@@ -766,7 +763,7 @@ def find_top_trace(xs) -> Trace:
                  default=None, key=attrgetter('level'))
   dynamic = thread_local_state.trace_state.trace_stack.dynamic
   top_main = (dynamic if top_main is None or dynamic.level > top_main.level
-                else top_main)
+              else top_main)
   return top_main and top_main.with_cur_sublevel()  # type: ignore
Rm old attribute annotations from TraceStack