# Buggy code
def _rewrite_eqn(
    eqn: core.JaxprEqn,
    eqns: List[core.JaxprEqn],
    input_token_var: core.Var,
    output_token_var: core.Var,
    mk_new_var: Callable[[core.AbstractValue], core.Var],
):
    """Rewrite an `eqn` and append equations to `eqns`.

    Assume that the current token is in `input_token_var` and the resulting
    token must end in `output_token_var`.
    """
    if eqn.primitive is id_tap_p:
        assert "has_token_" not in eqn.params
        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(eqn.params, has_token_=True),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.while_p:
        cond_jaxpr, _, body_jaxpr, _ = util.split_dict(
            eqn.params, ["cond_jaxpr", "cond_nconsts", "body_jaxpr", "body_nconsts"]
        )
        if xla.jaxpr_uses_outfeed(cond_jaxpr.jaxpr):
            _rewrite_while_outfeed_cond(
                eqn, eqns, input_token_var, output_token_var, mk_new_var
            )
            return

        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    body_jaxpr=_rewrite_typed_jaxpr(body_jaxpr, True, True),
                    cond_jaxpr=_rewrite_typed_jaxpr(cond_jaxpr, True, False),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.cond_p:
        branches, linear = util.split_dict(eqn.params, ["branches", "linear"])
        index, *operands = eqn.invars
        new_invars = [index, *operands, input_token_var]
        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    branches=tuple(
                        _rewrite_typed_jaxpr(jaxpr, True, True) for jaxpr in branches
                    ),
                    linear=(*linear, False),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.scan_p:
        num_consts, num_carry, carry_jaxpr, linear, _, _, _ = util.split_dict(
            eqn.params,
            [
                "num_consts",
                "num_carry",
                "jaxpr",
                "linear",
                "reverse",
                "length",
                "unroll",
            ],
        )
        # We add the token right at the end of carry
        nr_const_and_carry = num_consts + num_carry
        new_invars = (
            eqn.invars[0:nr_const_and_carry]
            + [input_token_var]
            + eqn.invars[nr_const_and_carry:]
        )
        new_jaxpr = _rewrite_typed_jaxpr(carry_jaxpr, True, True)
        # The rewrite has put the token at end, it has to be at end of carry
        new_jaxpr_invars = new_jaxpr.jaxpr.invars
        new_jaxpr_invars = (
            new_jaxpr_invars[0:nr_const_and_carry]
            + [new_jaxpr_invars[-1]]
            + new_jaxpr_invars[nr_const_and_carry:-1]
        )
        new_jaxpr.jaxpr.invars = new_jaxpr_invars
        new_jaxpr.in_avals = [v.aval for v in new_jaxpr_invars]

        new_jaxpr_outvars = new_jaxpr.jaxpr.outvars
        new_jaxpr_outvars = (
            new_jaxpr_outvars[0:num_carry]
            + [new_jaxpr_outvars[-1]]
            + new_jaxpr_outvars[num_carry:-1]
        )
        new_jaxpr.jaxpr.outvars = new_jaxpr_outvars
        new_jaxpr.out_avals = [v.aval for v in new_jaxpr_outvars]
        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                # Output token is at the end of carry result
                eqn.outvars[0:num_carry] + [output_token_var] + eqn.outvars[num_carry:],
                eqn.primitive,
                dict(
                    eqn.params,
                    jaxpr=new_jaxpr,
                    num_carry=num_carry + 1,
                    linear=linear + (False,),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is xla.xla_call_p:
        call_jaxpr = cast(core.Jaxpr, eqn.params["call_jaxpr"])
        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(eqn.params, call_jaxpr=_rewrite_jaxpr(call_jaxpr, True, True)),
                eqn.source_info,
            )
        )
    elif eqn.primitive is custom_derivatives.custom_jvp_call_jaxpr_p:
        fun_jaxpr = eqn.params["fun_jaxpr"]
        new_invars = [*eqn.invars, input_token_var]

        def unreachable_thunk():
            assert False, "Should not be reached"

        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
                    jvp_jaxpr_thunk=unreachable_thunk,
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is custom_derivatives.custom_vjp_call_jaxpr_p:
        fun_jaxpr = eqn.params["fun_jaxpr"]
        new_invars = [*eqn.invars, input_token_var]

        def unreachable_thunk():
            assert False, "Should not be reached"

        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
                    fwd_jaxpr_thunk=unreachable_thunk,
                    # The following are illegal values for the parameters, they
                    # should not be needed because this rewrite is just before
                    # compilation to XLA, which does not use those parameters.
                    bwd="illegal param",
                    out_trees="illegal param",
                ),
                eqn.source_info,
            )
        )
    else:
        raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")


# 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_eqn(
    eqn: core.JaxprEqn,
    eqns: List[core.JaxprEqn],
    input_token_var: core.Var,
    output_token_var: core.Var,
    mk_new_var: Callable[[core.AbstractValue], core.Var],
):
    """Rewrite an `eqn` and append equations to `eqns`.

    Assume that the current token is in `input_token_var` and the resulting
    token must end in `output_token_var`.
    """
    if eqn.primitive is id_tap_p:
        assert "has_token_" not in eqn.params
        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(eqn.params, has_token_=True),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.while_p:
        cond_jaxpr, _, body_jaxpr, _ = util.split_dict(
            eqn.params, ["cond_jaxpr", "cond_nconsts", "body_jaxpr", "body_nconsts"]
        )
        if xla.jaxpr_uses_outfeed(cond_jaxpr.jaxpr):
            _rewrite_while_outfeed_cond(
                eqn, eqns, input_token_var, output_token_var, mk_new_var
            )
            return

        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    body_jaxpr=_rewrite_typed_jaxpr(body_jaxpr, True, True),
                    cond_jaxpr=_rewrite_typed_jaxpr(cond_jaxpr, True, False),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.cond_p:
        branches, linear = util.split_dict(eqn.params, ["branches", "linear"])
        index, *operands = eqn.invars
        new_invars = [index, *operands, input_token_var]
        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    branches=tuple(
                        _rewrite_typed_jaxpr(jaxpr, True, True) for jaxpr in branches
                    ),
                    linear=(*linear, False),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is lax.scan_p:
        num_consts, num_carry, carry_jaxpr, linear, _, _, _ = util.split_dict(
            eqn.params,
            [
                "num_consts",
                "num_carry",
                "jaxpr",
                "linear",
                "reverse",
                "length",
                "unroll",
            ],
        )
        # We add the token right at the end of carry
        nr_const_and_carry = num_consts + num_carry
        new_invars = (
            eqn.invars[0:nr_const_and_carry]
            + [input_token_var]
            + eqn.invars[nr_const_and_carry:]
        )
        new_jaxpr = _rewrite_typed_jaxpr(carry_jaxpr, True, True)
        # The rewrite has put the token at end, it has to be at end of carry
        new_jaxpr_invars = new_jaxpr.jaxpr.invars
        new_jaxpr_invars = (
            new_jaxpr_invars[0:nr_const_and_carry]
            + [new_jaxpr_invars[-1]]
            + new_jaxpr_invars[nr_const_and_carry:-1]
        )
        new_jaxpr.jaxpr.invars = new_jaxpr_invars
        new_jaxpr.in_avals = [v.aval for v in new_jaxpr_invars]

        new_jaxpr_outvars = new_jaxpr.jaxpr.outvars
        new_jaxpr_outvars = (
            new_jaxpr_outvars[0:num_carry]
            + [new_jaxpr_outvars[-1]]
            + new_jaxpr_outvars[num_carry:-1]
        )
        new_jaxpr.jaxpr.outvars = new_jaxpr_outvars
        new_jaxpr.out_avals = [v.aval for v in new_jaxpr_outvars]
        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                # Output token is at the end of carry result
                eqn.outvars[0:num_carry] + [output_token_var] + eqn.outvars[num_carry:],
                eqn.primitive,
                dict(
                    eqn.params,
                    jaxpr=new_jaxpr,
                    num_carry=num_carry + 1,
                    linear=linear + (False,),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is xla.xla_call_p:
        call_jaxpr = cast(core.Jaxpr, eqn.params["call_jaxpr"])
        eqns.append(
            core.new_jaxpr_eqn(
                eqn.invars + [input_token_var],
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    call_jaxpr=_rewrite_jaxpr(call_jaxpr, True, True),
                    donated_invars=eqn.params["donated_invars"] + (False,),
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is custom_derivatives.custom_jvp_call_jaxpr_p:
        fun_jaxpr = eqn.params["fun_jaxpr"]
        new_invars = [*eqn.invars, input_token_var]

        def unreachable_thunk():
            assert False, "Should not be reached"

        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
                    jvp_jaxpr_thunk=unreachable_thunk,
                ),
                eqn.source_info,
            )
        )
    elif eqn.primitive is custom_derivatives.custom_vjp_call_jaxpr_p:
        fun_jaxpr = eqn.params["fun_jaxpr"]
        new_invars = [*eqn.invars, input_token_var]

        def unreachable_thunk():
            assert False, "Should not be reached"

        eqns.append(
            core.new_jaxpr_eqn(
                new_invars,
                eqn.outvars + [output_token_var],
                eqn.primitive,
                dict(
                    eqn.params,
                    fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
                    fwd_jaxpr_thunk=unreachable_thunk,
                    # The following are illegal values for the parameters, they
                    # should not be needed because this rewrite is just before
                    # compilation to XLA, which does not use those parameters.
                    bwd="illegal param",
                    out_trees="illegal param",
                ),
                eqn.source_info,
            )
        )
    else:
        raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")