p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
            carry_causal = init_carry_causal
            for step_n in range(N):
                carry_causal, _ = causal_step_fn(
                    carry_causal, jnp.array(step_n, dtype=jnp.int32)
                )
            final_carry_causal = carry_causal