def get_conds(fun):
    jaxpr = api.make_jaxpr(api.grad(fun))(0.0, 0)
    return [eqn for eqn in jaxpr.jaxpr.eqns if (eqn.primitive.name == 'cond')]
def get_conds(fun):
    jaxpr = api.make_jaxpr(api.grad(fun))(0.0, 0)
    return [eqn for eqn in jaxpr.jaxpr.eqns if (eqn.primitive.name == 'cond')]