def branch_outvars_len(cond_eqn):
    lens = [len(jaxpr.jaxpr.outvars) for jaxpr in cond_eqn.params['branches']]
    assert (len(set(lens)) == 1)
    return lens[0]
def branch_outvars_len(cond_eqn):
    lens = [len(jaxpr.jaxpr.outvars) for jaxpr in cond_eqn.params['branches']]
    assert (len(set(lens)) == 1)
    return lens[0]