def new_jaxpr():
    jaxpr = api.make_jaxpr((lambda x: lax.switch(0, [jnp.sin, jnp.cos], x)))(1.0).jaxpr
    cond = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'cond')))
    return (jaxpr, cond)
def new_jaxpr():
    jaxpr = api.make_jaxpr((lambda x: lax.switch(0, [jnp.sin, jnp.cos], x)))(1.0).jaxpr
    cond = next((eqn for eqn in jaxpr.eqns if (eqn.primitive.name == 'cond')))
    return (jaxpr, cond)