Independent vector-Jacobian product for multiple arguments

[chaoming0625]

In the official documentation, JAX supports vjp for multiple arguments:

>>> import jax
>>>
>>> def f(x, y): return jax.numpy.sin(x), jax.numpy.cos(y)
>>> primals, f_vjp = jax.vjp(f, 0.5, 1.0)
>>> xbar, ybar = f_vjp((-0.7, 0.3))
>>> print(xbar)
-0.61430776
>>> print(ybar)
-0.2524413

Once the return values are using both arguments:

import jax.numpy as jnp
from jax import vjp

def vgrad(f, xs):
  ys, vjp_fn = vjp(f, *xs)
  return vjp_fn(tuple([jnp.ones(y.shape) for y in ys]))


def f(x, y):
  dx = x ** 2 + y ** 2 + 10
  dy = x ** 3 + y ** 3 - 10
  return dx, dy

pprint(vgrad(f, [0.1, 0.2]))

# (0.23, 0.52)

However, we want to get the derivative for each argument independently. That is to say, the except derivatives are:

def expect(x, y):
  return 2 * x, 3 * y ** 2

expect(0.1, 0.2)

# (0.2, 0.12)

Can anyone tell me how to make the excepted gradients? Thanks!

Note here we cannot use custom_vjp because the do not explicitly know the logic of the provided functions.

[chaoming0625]

Can anyone help me? Thanks!

[hawkinsp]

I'm having difficulty understanding what you want. I don't follow the question about expect. Let's forget the Python code for the moment: what do you want to compute mathematically?

If you want the gradient of a function with respect to one of the arguments, you might want argnums argument to grad (https://jax.readthedocs.io/en/latest/jax.html?highlight=grad#jax.grad). If you want something else, you'll have to explain better what you are trying to compute.

[chaoming0625]

Dear @hawkinsp , thanks for your reply! I appologize for my unclear question.

Actually, I have three equations:

Hope you can see the image, If not, there are the latex codes:

$$ \begin{aligned} f1(V, m, h) = &-\bar{g}_{Na} m^3 h (V -E_{Na})\\ f2(V, m) = &\frac{1}{1-\exp(\frac{-(V + 40)} {10})} * m + 4.0 \exp \left(\frac{-(V+65)}{18}\right) (1- m) \\ f3(V, h) = &0.07 \exp(\frac{-(V+65)}{20}) * h + \frac{1}{1+\exp \left(\frac{-(V+35)}{10}\right)} (1 -h) \end{aligned} $$

I am trying to get the derivative: df1/dV, df2/dm and df3/dh. So I make the following codes:

from functools import partial
import jax.numpy as jnp
from jax import random
from jax import vjp

# Vector-valued gradients with VJPs
# https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#vector-valued-gradients-with-vjps
def vgrad(f, xs):
  ys, vjp_fn = vjp(f, *xs)
  return vjp_fn(tuple([jnp.ones(y.shape) for y in ys]))

ENa = 55.
gNa = 35.


def f1(V, m, h):
  return - gNa * m ** 3 * h * (V - ENa),


def f2(m, V):
  a = 1 / (1 - jnp.exp(-(V + 40) / 10))
  b = 4 * jnp.exp(-(V + 65) / 18)
  return a * m + b * (1 - m),


def f3(h, V):
  a = 0.07 * jnp.exp(-(V + 65) / 20)
  b = 1 / (1 + jnp.exp(-(V + 35) / 10))
  return a * h + b * (1 - h),


def F(V, m, h):
  f1v = f1(V, m, h)[0]
  f2v = f2(m, V)[0]
  f3v = f3(h, V)[0]
  return f1v, f2v, f3v

When I simultaneously evaluate the function f1, f2 and f3 by the merged function F. I got the following gradients:

key = random.PRNGKey(0)
key, k1, k2, k3 = random.split(key, 4)
ex_V = random.uniform(k1, (3,))
ex_m = random.uniform(k2, (3,))
ex_h = random.uniform(k3, (3,))

D = vgrad(F, (ex_V, ex_m, ex_h))
print(D[0])
# [-0.4800583 -9.6292515 -4.302372 ]

print(D[1])
# [ 249.06578 2387.181    911.52954]

print(D[2])
# [ 59.27976 564.16797 904.1445 ]

However, what I want to are the gradient values of:

print(vgrad(partial(f1, m=ex_m, h=ex_h), (ex_V,))[0])
# [-0.47693872 -9.626045   -4.301631  ]

print(vgrad(partial(f2, V=ex_V), (ex_m,))[0])
# [0.9116234  0.9108186  0.91096807]

print(vgrad(partial(f3, V=ex_V), (ex_h,))[0])
# [-0.9687328  -0.96815366 -0.9682618 ]

Obviously, the gradients are significantly different by using the separate and merged functions.

There are many reasons I must code these three functions f1, f2 and f3 into a merged function F. But I am also trying to get the vector gradients for indepedent input/output mappings. Sometimes, the merged function is somehow complex. For example, the size of the argument may be different. This result in vmap is not a good idea to apply to this situation.

Could you give me a solution? Thank you very very much!!

[jakevdp]

I think if you want three independent gradients, you'll have to compute the three gradients independently. Your last few lines of code (calling vgrad three times with three different partials) is precisely how I'd do it.

[chaoming0625]

Dear @jakevdp , I am so much grateful for your reply. Thanks.