def f(x):
    return cond((x < 2), (lambda x: (np.array([1.0, 2.0]) * x)), (lambda x: (np.array([3.0, 4.0]) * jnp.sin(x))), x)
def f(x):
    return cond((x < 2), (lambda x: (np.array([1.0, 2.0]) * x)), (lambda x: (np.array([3.0, 4.0]) * jnp.sin(x))), x)