A naive question about the extend vjp used in netket.

[yiminghwang]

Hi, there,

I have a naive question when reading the source codes of extend vjp in netket,
Because Jax does not support $f:\mathbb{R}\rightarrow\mathbb{C}$, we cannot directly use it to get the gradient of the logarithmic complex wave function generated by NN, such as $\log(\Psi(\sigma))=A_{\theta_r}(\sigma)+jB_{\theta_i}(\sigma), \theta_r,\theta_i \in \mathbb{R}$. thus, we have to extend the original vjp function in Jax.
I follow the source codes in _vjp.py to deal with the above situation.

def vjp_rc(
    fun: Callable, *primals, has_aux: bool = False, conjugate: bool = False
) -> Union[Tuple[Any, Callable], Tuple[Any, Callable, Any]]:
    if has_aux:

        def real_fun(*primals):
            val, aux = fun(*primals)
            return val.real, aux

        def imag_fun(*primals):
            val, aux = fun(*primals)
            return val.imag, aux

        vals_r, vjp_r_fun, aux = jax.vjp(real_fun, *primals, has_aux=True)
        vals_j, vjp_j_fun, _ = jax.vjp(imag_fun, *primals, has_aux=True)

    else:
        real_fun = lambda *primals: fun(*primals).real
        imag_fun = lambda *primals: fun(*primals).imag

        vals_r, vjp_r_fun = jax.vjp(real_fun, *primals, has_aux=False)
        vals_j, vjp_j_fun = jax.vjp(imag_fun, *primals, has_aux=False)

    primals_out = vals_r + 1j * vals_j

    def vjp_fun(ȳ):
        """
        function computing the vjp product for a R->C function.
        """
        ȳ_r = ȳ.real
        ȳ_j = ȳ.imag

        # val = vals_r + vals_j
        vr_jr = vjp_r_fun(jnp.asarray(ȳ_r, dtype=vals_r.dtype))
        vj_jr = vjp_r_fun(jnp.asarray(ȳ_j, dtype=vals_r.dtype))
        vr_jj = vjp_j_fun(jnp.asarray(ȳ_r, dtype=vals_j.dtype))
        vj_jj = vjp_j_fun(jnp.asarray(ȳ_j, dtype=vals_j.dtype))

        r = tree_map(
            lambda re, im: re + 1j * im,
            vr_jr,
            vj_jr,
        )
        i = tree_map(lambda re, im: re + 1j * im, vr_jj, vj_jj)
        out = tree_map(lambda re, im: re + 1j * im, r, i)

        if conjugate:
            out = tree_map(jnp.conjugate, out)

        return out

    if has_aux:
        return primals_out, vjp_fun, aux
    else:
        return primals_out, vjp_fun

As we all known, the jax.vjp is used for calculating the gradient in reverse-mode, such as

    def f(x, y):
        return x+1, 2*y**2-1

    primals, f_vjp = jax.vjp(f, 1., 2.)
    xbar, ybar = f_vjp((1.,1.))
    print(xbar)
    print(ybar)

we calculate the gradient of x and y when the inputs are x=1. and y=2.

my question is why it returns vjp_fun at the end instead of [vjp_r_fun,vjp_i_fun]?

P.S.
For testing, I write the following codes,

  def f(x, y):
          return x**2+2*j*y

    params = jnp.array([1., 2.])
    _, vjp_funs = vjp(lambda w: func(w),
                      params,
                      conjugate=True) # if I return [vjp_r_fun,vjp_i_fun]

    grad_r = vjp_funs[0](1.)
    grad_i = vjp_funs[1](1.)
    grad = grad_r+grad_i

it returns me grad_r = [2., 2.], it seems right when we need to separately update the real parameters x and y.
when I try to use vjp_fun as the return, i.e. grad = vjp_fun(_), it returns me [2.-8j, -8-2.j].
I am curious whether I missing something or using it in the wrong way.

[PhilipVinc]

I don't understand what you are saying, also because your snippets don't run for me.
Can you write a snippet that does not run and maybe the mathematical formula corresponding to it so that we can see what is wrong?

For info, nk.jax.vjp(f, z)(v) computes the following, assuming $f(z) = u(x, y) + iv(x, y) $, $z =x+iy $ and vector $v=c+id$ .

I suspect that whatever you might be computing is not that.

[yiminghwang]

Thanks for your quick answer!
sorry about the misunderstanding discription. let me make myself clear,
I want to calculate the gradient of a function $f:R \rightarrow C$ , and $f=f_A(\theta_a,x)+1j*f_B(\theta_b,x)$, $f_A$ and $f_B$ only contains real parameters $\theta_a$ and $\theta_b$, and $x$ is the input. Actually, you can regard the $f_A$ and $f_B$ are two independent components in $f$.
our goal is to calculate the gradient of $f$, i.e. $\partial_{\theta_a}{f}$ and $\partial_{\theta_b}{f}$,
As Jax does not support $f: R\rightarrow C$(parameters are real, $f$ is complex), thus we need extend the vjp method of Jax.
In netket, the file _vjp.py contains the code for $f: R\rightarrow C$, so I run the following codes,

    from netket.jax._vjp import vjp

    def func(params):
        return params[0] ** 2 + 2 * 1j * params[1]

    params = jnp.array([1., 2.]) # we restrict the parameters are all real
    _, vjp_funs = vjp(lambda w: func(w),
                      params,
                      conjugate=True)

    grad = vjp_funs(_)
    print(grad)

it returns the grad: [2.-8j, -8-2.j],
my question is why the grad is not only real? am I misunderstanding about the usage of vjp in netket?

[PhilipVinc]

I think you are misunderstanding what vjp means.
vjp != grad!

For a $R\rightarrow C$ function the gradient is

$$ \partial_{\theta_a} f = 1 \partial_{\theta_a} f $$

By which I mean that you are multiplying the Jacobian with a vector containing all 1s. Note that this is not true in general (look at the formula I wrote above). It's only true for this particular case (and to an extent for holomorphic functions).

So, the problem is that you are feeding to vjp_funs your previous result _ is

In [6]: _
Out[8]: DeviceArray(1.+4.j, dtype=complex128)

so you are not computing the gradient, you are computing something else. You are actually computing grad -4j grad or something.
If you were to actually compute the gradient, you'd see you get the correct object

In [9]: vjp_funs(1+4j)
Out[9]: (DeviceArray([ 2.-8.j, -8.-2.j], dtype=complex128),)

In [10]: vjp_funs(1)
Out[10]: (DeviceArray([2.-0.j, 0.-2.j], dtype=complex128),)

[yiminghwang]

get it, thank you so much!