can we init the l_out inside the mnist forward and returns it? This way the only additional output we provide is the count.

thanks, good idea, due to hard to gain result before finish one run, so current plan would be:
run body_fn with pure python while, get result, then use this result as l_out

before implementation, did quick benchmark for this idea with code:

- iteri = 3: 
  while_loop(3.675 s), pure while(3.107 s), pure_while_then_while_loop_plan(3.724 s)
- iteri = 10: 
  while_loop(3.670s), pure while(3.119s), pure_while_then_while_loop_plan(3.932s)
- iteri = 10: 
  while_loop(3.744 s), pure while(3.110 s), pure_while_then_while_loop_plan(3.749 s)
- iteri = 10: 
  while_loop(3.878 s), pure while(7.053 s), pure_while_then_while_loop_plan(3.809 s)

so this plan looks like worth to try