tweak

Author: mattjj

jax/_src/lax/control_flow/loops.py

@@ -872,11 +872,11 @@ def _scan_partial_eval(trace, *tracers, reverse: bool,
   # At this point all non-forwarded residuals are treated as extensive outputs
   # of jaxpr_known. Hoist out those that only depend on consts.
   #   Before: jaxpr_known: [*known_ins] -> [*known_outs, *non_fwd_res]
-  #   After: jaxpr_known: [*known_consts, *known_ins] -> [*known_outs, *ext_res]
+  #   After: jaxpr_known: [*known_consts_, *known_ins] -> [*known_outs, *ext_res]
   # where, modulo hoisted res not being broadcast, we have
   #   non_fwd_res = merge_lists(which_hoisted, ext_res, hoisted_res)
   known_consts, known_ins = split_list(known_ins, [num_consts_known])
-  jaxpr_known_, known_consts_, which_hoisted, hoisted_res = \
+  jaxpr_known, known_consts_, which_hoisted, hoisted_res = \
       _scan_known_hoisting(jaxpr_known, known_consts, num_res_out)
 
   # To make jaxpr_unknown match the scan calling convention, move to the back
@@ -894,13 +894,13 @@ def _scan_partial_eval(trace, *tracers, reverse: bool,
 
   # Run the known part of the scan (if it has any outputs or effects).
   linear_known, linear_unknown = partition_list(unknowns, linear)
-  if not jaxpr_known_.out_avals and not jaxpr_known_.effects:
+  if not jaxpr_known.out_avals and not jaxpr_known.effects:
     known_outs_ext_res = []
   else:
-    linear_known = [False] * len(jaxpr_known_.in_avals)  # TODO conservative
-    assert len(known_consts_) + len(known_ins) == len(jaxpr_known_.in_avals)
+    linear_known = [False] * len(jaxpr_known.in_avals)  # TODO conservative
+    assert len(known_consts_) + len(known_ins) == len(jaxpr_known.in_avals)
     known_outs_ext_res = scan_p.bind(
-        *known_consts_, *known_ins, jaxpr=jaxpr_known_, reverse=reverse,
+        *known_consts_, *known_ins, jaxpr=jaxpr_known, reverse=reverse,
         length=length, num_consts=len(known_consts_),
         num_carry=num_carry_known, linear=(*linear_known,), unroll=unroll,
         _split_transpose=_split_transpose)