CODEWITHSUNDEEP
pythonjaxautodiff
user654123
user654123

Reputation: 485

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

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, i.e.,

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

but that does not work.

Upvotes: 2

Views: 1848

Answers (1)

jakevdp
jakevdp

Reputation: 86513

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.]

Upvotes: 3

Related Questions