def _run_loop_in_debug(cond_fn, body_fn, init_state):
        """
        Run generation in untraced mode. This should only be used for debugging purposes.
        """
        state = init_state
        while cond_fn(state):
            state = body_fn(state)
        return state

        Run generation in untraced mode. This should only be used for debugging purposes.
        
dpkd/transformers/src/transformers/generation/flax_utils.py