Jax: Take derivative with respect to index of vector-valued argument

Question

Does Jax support taking the derivate w.r.t. an index of a vector-valued variable? Consider this example (where a is a vector/array):

def test_func(a):
  return a[0]**a[1]

I can pass in the argument number into grad(..) , but I cannot seem to pass the index of a vector-valued argument like in the example above. I tried passing a tuple of tuples, ie,

grad(test_func, argnums=((0,),))

but that does not work.

Answer 1

There's no built-in transform that can take gradients with respect to certain elements of arrays, but you can straightforwardly do this via a wrapper function that splits the array into individual elements; for example:

import jax
import jax.numpy as jnp

def test_func(a):
  return a[0]**a[1]

a = jnp.array([1.0, 2.0])
fgrad = jax.grad(lambda *args: test_func(jnp.array(args)), argnums=0)
print(fgrad(*a))
# 2.0

If you want to take a gradient with respect to all the inputs individually (returning a vector of gradients with respect to each entry), you can use jax.jacobian :

print(jax.jacobian(test_func)(a))
# [2. 0.]