def fun_vjp_jax(args_jax, out_cts_jax):
    (_, pullback_jax) = jax.vjp(fun, *args_jax)
    return pullback_jax(out_cts_jax)
def fun_vjp_jax(args_jax, out_cts_jax):
    (_, pullback_jax) = jax.vjp(fun, *args_jax)
    return pullback_jax(out_cts_jax)