def METHOD_NAME(
    fun: Callable, *primals, has_aux: bool = False, conjugate: bool = False
) -> Union[tuple[Any, Callable], tuple[Any, Callable, Any]]:
    if has_aux:
        out, _vjp_fun, aux = jax.vjp(fun, *primals, has_aux=True)
    else:
        out, _vjp_fun = jax.vjp(fun, *primals, has_aux=False)
    def vjp_fun(ȳ):
        ȳ = jnp.asarray(ȳ, dtype=out.dtype)
        dȳ = _vjp_fun(ȳ)
        if conjugate:
            dȳ = tree_map(jnp.conjugate, dȳ)
        return dȳ
    if has_aux:
        return out, vjp_fun, aux
    else:
        return out, vjp_fun