def f_ref(x):
    z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else jnp.sin(x))
    return z.sum()
def f_ref(x):
    z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else jnp.sin(x))
    return z.sum()