def fun(x, y):
    return cond((x < 3), (), (lambda _: (2.0 * jnp.sin(y))), x, (lambda x: (2.0 * x)))
def fun(x, y):
    return cond((x < 3), (), (lambda _: (2.0 * jnp.sin(y))), x, (lambda x: (2.0 * x)))