def METHOD_NAME(forward_fn, params, samples, v, pdf=None):
    w = _O_jvp(forward_fn, params, samples, v)
    if pdf is None:
        w = w * (1.0 / (samples.shape[0] * samples.shape[1] * mpi.n_nodes))
        w_, chunk_fn = unchunk(w)
        w = chunk_fn(subtract_mean(w_))  # w/ MPI
    else:
        w_, chunk_fn = unchunk(w)
        # here we assume pdf is chunked,
        # which we undo as the flat vector is needed here
        pdf_, _ = unchunk(pdf)
        w_ = pdf_ * (w_ - mpi.mpi_sum_jax(pdf_ @ w_)[0])
        w = chunk_fn(w_)
    res = _OH_w(forward_fn, params, samples, w)
    return jax.tree_map(lambda x: mpi.mpi_sum_jax(x)[0], res)  # MPI