def f(x):
    return cond((x[0] < 2), (lambda x: (jnp.array([1.0, 2.0]) * x)), (lambda x: jnp.cos(jnp.sin(x))), x).sum()
def f(x):
    return cond((x[0] < 2), (lambda x: (jnp.array([1.0, 2.0]) * x)), (lambda x: jnp.cos(jnp.sin(x))), x).sum()