# Buggy code
def _rewrite_while_outfeed_cond(
    eqn: core.JaxprEqn,
    eqns: List[core.JaxprEqn],
    input_token_var: core.Var,
    output_token_var: core.Var,
    mk_new_var: Callable,
):
    """Rewrite a while whose cond has outfeed"""
    cond_jaxpr, cond_nconsts, body_jaxpr, body_nconsts = util.split_dict(
        eqn.params, ["cond_jaxpr", "cond_nconsts", "body_jaxpr", "body_nconsts"]
    )
    transformed_cond_jaxpr = _rewrite_typed_jaxpr(cond_jaxpr, True, True)
    carry_invars = eqn.invars[cond_nconsts + body_nconsts :]
    # pred1, token1 = rewrite(COND)(cond_consts, carry_invars, input_token)
    pred1_and_token1 = [
        mk_new_var(ov.aval) for ov in transformed_cond_jaxpr.jaxpr.outvars
    ]
    eqns.append(
        core.new_jaxpr_eqn(
            eqn.invars[0:cond_nconsts] + carry_invars + [input_token_var],
            pred1_and_token1,
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_cond_jaxpr.jaxpr,
                name="cond_before",
                donated_invars=(False,) * (cond_nconsts + len(carry_invars) + 1),
            ),
            eqn.source_info,
        )
    )
    # Make a new cond "lambda pred, carry, token: pred"
    new_cond_pred_invar = mk_new_var(cond_jaxpr.out_avals[0])
    new_cond_invars = (
        [new_cond_pred_invar]
        + [mk_new_var(cv.aval) for cv in carry_invars]
        + [mk_new_var(core.abstract_token)]
    )
    new_cond_jaxpr = _mk_typed_jaxpr(
        core.Jaxpr([], new_cond_invars, [new_cond_pred_invar], []), []
    )
    # Make a new body:
    #   "lambda cond_constvars, body_constvars, pred, carry, token:
    #        carry2, token2 = rewrite(BODY)(body_constvars, carry, token)
    #        pred2, token3 = rewrite(COND)(cond_constvars, carry2, token2)
    #        (pred2, carry2, token3)
    transformed_body_jaxpr = _rewrite_typed_jaxpr(body_jaxpr, True, True)
    new_body_invars_cond_constvars = [
        mk_new_var(v.aval) for v in eqn.invars[0:cond_nconsts]
    ]
    new_body_invars_body_constvars = [
        mk_new_var(v.aval)
        for v in eqn.invars[cond_nconsts : cond_nconsts + body_nconsts]
    ]
    new_body_invars_pred = mk_new_var(cond_jaxpr.out_avals[0])
    new_body_invars_carry = [mk_new_var(cv.aval) for cv in carry_invars]
    new_body_invars_token = mk_new_var(core.abstract_token)

    new_body_carry2 = [mk_new_var(cv.aval) for cv in carry_invars]
    new_body_token2 = mk_new_var(core.abstract_token)
    new_body_pred2 = mk_new_var(cond_jaxpr.out_avals[0])
    new_body_token3 = mk_new_var(core.abstract_token)

    new_body_eqns = [
        core.new_jaxpr_eqn(
            new_body_invars_body_constvars
            + new_body_invars_carry
            + [new_body_invars_token],
            new_body_carry2 + [new_body_token2],
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_body_jaxpr.jaxpr,
                name="body",
                donated_invars=(False,)
                * (
                    len(new_body_invars_body_constvars)
                    + len(new_body_invars_carry)
                    + 1
                    + len(new_body_carry2)
                    + 1
                ),
            ),
            eqn.source_info,
        ),
        core.new_jaxpr_eqn(
            new_body_invars_cond_constvars + new_body_carry2 + [new_body_token2],
            [new_body_pred2, new_body_token3],
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_cond_jaxpr.jaxpr,
                name="cond_body",
                donated_invars=(False,)
                * (len(new_body_invars_cond_constvars) + len(new_body_carry2) + 1 + 2),
            ),
            eqn.source_info,
        ),
    ]
    new_body_jaxpr = _mk_typed_jaxpr(
        core.Jaxpr(
            [],
            (
                new_body_invars_cond_constvars
                + new_body_invars_body_constvars
                + [new_body_invars_pred]
                + new_body_invars_carry
                + [new_body_invars_token]
            ),
            ([new_body_pred2] + new_body_carry2 + [new_body_token3]),
            new_body_eqns,
        ),
        [],
    )

    pred_out = mk_new_var(cond_jaxpr.out_avals[0])
    eqns.append(
        core.new_jaxpr_eqn(
            (
                eqn.invars[0 : cond_nconsts + body_nconsts]
                + [pred1_and_token1[0]]
                + carry_invars
                + [pred1_and_token1[1]]
            ),
            ([pred_out] + eqn.outvars + [output_token_var]),
            lax.while_p,
            dict(
                cond_jaxpr=new_cond_jaxpr,
                cond_nconsts=0,
                body_jaxpr=new_body_jaxpr,
                body_nconsts=cond_nconsts + body_nconsts,
            ),
            eqn.source_info,
        )
    )


# Traceback (AssertionError)
Traceback (most recent call last):
File "demos/encoding.py", line 95, in <module>
encoding_demo()
File "demos/encoding.py", line 80, in encoding_demo
training_trajectory = problem.train()
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/cmm/problem/problem.py", line 31, in train
trajectory, self.model_parameters, self.iterated_function = train_model(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/cmm/model/inference.py", line 76, in train_model
augmented, trajectory = method(model_meta_parameters,
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/api.py", line 169, in f_jitted
out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/core.py", line 1103, in call_bind
outs = primitive.impl(fun, *args, **params)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/xla.py", line 532, in _xla_call_impl
compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/linear_util.py", line 221, in memoized_fun
ans = call(fun, *args)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/xla.py", line 641, in _xla_callable
out_nodes = jaxpr_subcomp(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/xla.py", line 410, in jaxpr_subcomp
ans = rule(c, axis_env, extend_name_stack(name_stack, eqn.primitive.name),
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/xla.py", line 912, in f
jaxpr, _, consts = pe.trace_to_jaxpr(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 429, in trace_to_jaxpr
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/linear_util.py", line 150, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 1319, in _scan_impl
return _scan_impl_loop(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 1280, in _scan_impl_loop
_, *outs = while_loop(cond_fun, body_fun, init_val)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 277, in while_loop
body_jaxpr, body_consts, body_tree = _initial_style_jaxpr(body_fun, in_tree, init_avals)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 72, in _initial_style_jaxpr
jaxpr, out_pvals, consts, out_tree = _initial_style_untyped_jaxpr(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 66, in _initial_style_untyped_jaxpr
jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 429, in trace_to_jaxpr
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/linear_util.py", line 150, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/lax/lax_control_flow.py", line 1270, in body_fun
out_flat = f_impl(*consts, *carry, *x)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/core.py", line 142, in jaxpr_as_fun
return eval_jaxpr(typed_jaxpr.jaxpr, typed_jaxpr.literals, *args)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/core.py", line 346, in eval_jaxpr
ans = eqn.primitive.bind(*(subfuns + in_vals), **params)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/core.py", line 1106, in call_bind
outs = primitive.process(top_trace, fun, tracers, params)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/core.py", line 1115, in process
return trace.process_call(self, fun, tracers, params)
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 232, in process_call
new_params = update_params(new_params, [not t.pval.is_known() for t in tracers])
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/interpreters/xla.py", line 813, in _xla_call_partial_eval_update_params
donated_invars = [d for d, uk in zip(donated_invars, in_unknowns) if uk]
File "/home/neil/.pyenv/versions/3.8.2/lib/python3.8/site-packages/jax/util.py", line 26, in safe_zip
assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
AssertionError: length mismatch: [18, 19]

# Fixed code
def _rewrite_while_outfeed_cond(
    eqn: core.JaxprEqn,
    eqns: List[core.JaxprEqn],
    input_token_var: core.Var,
    output_token_var: core.Var,
    mk_new_var: Callable,
):
    """Rewrite a while whose cond has outfeed"""
    cond_jaxpr, cond_nconsts, body_jaxpr, body_nconsts = util.split_dict(
        eqn.params, ["cond_jaxpr", "cond_nconsts", "body_jaxpr", "body_nconsts"]
    )
    transformed_cond_jaxpr = _rewrite_typed_jaxpr(cond_jaxpr, True, True)
    carry_invars = eqn.invars[cond_nconsts + body_nconsts :]
    # pred1, token1 = rewrite(COND)(cond_consts, carry_invars, input_token)
    pred1_and_token1 = [
        mk_new_var(ov.aval) for ov in transformed_cond_jaxpr.jaxpr.outvars
    ]
    eqns.append(
        core.new_jaxpr_eqn(
            eqn.invars[0:cond_nconsts] + carry_invars + [input_token_var],
            pred1_and_token1,
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_cond_jaxpr.jaxpr,
                name="cond_before",
                donated_invars=(False,) * len(transformed_cond_jaxpr.in_avals),
            ),
            eqn.source_info,
        )
    )
    # Make a new cond "lambda pred, carry, token: pred"
    new_cond_pred_invar = mk_new_var(cond_jaxpr.out_avals[0])
    new_cond_invars = (
        [new_cond_pred_invar]
        + [mk_new_var(cv.aval) for cv in carry_invars]
        + [mk_new_var(core.abstract_token)]
    )
    new_cond_jaxpr = _mk_typed_jaxpr(
        core.Jaxpr([], new_cond_invars, [new_cond_pred_invar], []), []
    )
    # Make a new body:
    #   "lambda cond_constvars, body_constvars, pred, carry, token:
    #        carry2, token2 = rewrite(BODY)(body_constvars, carry, token)
    #        pred2, token3 = rewrite(COND)(cond_constvars, carry2, token2)
    #        (pred2, carry2, token3)
    transformed_body_jaxpr = _rewrite_typed_jaxpr(body_jaxpr, True, True)
    new_body_invars_cond_constvars = [
        mk_new_var(v.aval) for v in eqn.invars[0:cond_nconsts]
    ]
    new_body_invars_body_constvars = [
        mk_new_var(v.aval)
        for v in eqn.invars[cond_nconsts : cond_nconsts + body_nconsts]
    ]
    new_body_invars_pred = mk_new_var(cond_jaxpr.out_avals[0])
    new_body_invars_carry = [mk_new_var(cv.aval) for cv in carry_invars]
    new_body_invars_token = mk_new_var(core.abstract_token)

    new_body_carry2 = [mk_new_var(cv.aval) for cv in carry_invars]
    new_body_token2 = mk_new_var(core.abstract_token)
    new_body_pred2 = mk_new_var(cond_jaxpr.out_avals[0])
    new_body_token3 = mk_new_var(core.abstract_token)

    new_body_eqns = [
        core.new_jaxpr_eqn(
            new_body_invars_body_constvars
            + new_body_invars_carry
            + [new_body_invars_token],
            new_body_carry2 + [new_body_token2],
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_body_jaxpr.jaxpr,
                name="body",
                donated_invars=(False,) * len(transformed_body_jaxpr.in_avals),
            ),
            eqn.source_info,
        ),
        core.new_jaxpr_eqn(
            new_body_invars_cond_constvars + new_body_carry2 + [new_body_token2],
            [new_body_pred2, new_body_token3],
            xla.xla_call_p,
            dict(
                call_jaxpr=transformed_cond_jaxpr.jaxpr,
                name="cond_body",
                donated_invars=(False,) * len(transformed_cond_jaxpr.in_avals),
            ),
            eqn.source_info,
        ),
    ]
    new_body_jaxpr = _mk_typed_jaxpr(
        core.Jaxpr(
            [],
            (
                new_body_invars_cond_constvars
                + new_body_invars_body_constvars
                + [new_body_invars_pred]
                + new_body_invars_carry
                + [new_body_invars_token]
            ),
            ([new_body_pred2] + new_body_carry2 + [new_body_token3]),
            new_body_eqns,
        ),
        [],
    )

    pred_out = mk_new_var(cond_jaxpr.out_avals[0])
    eqns.append(
        core.new_jaxpr_eqn(
            (
                eqn.invars[0 : cond_nconsts + body_nconsts]
                + [pred1_and_token1[0]]
                + carry_invars
                + [pred1_and_token1[1]]
            ),
            ([pred_out] + eqn.outvars + [output_token_var]),
            lax.while_p,
            dict(
                cond_jaxpr=new_cond_jaxpr,
                cond_nconsts=0,
                body_jaxpr=new_body_jaxpr,
                body_nconsts=cond_nconsts + body_nconsts,
            ),
            eqn.source_info,
        )
    )