How to calculate the partial derivative in an efficient way

[wuwudawen]

Consider a function with two inputs, $f(x, y): \mathbb{R}^2 \to \mathbb{R}$.
I want to calculate the partial derivative with respect to $x$, i.e., $\frac{\partial f}{\partial x}$, at a certain point $(x, y)$.

Using grad(f) and select the first element can definitely do the job, but is there a more efficient way?
Since grad(f) calculate the whole gradient vector, but what I want is one value of it, and I don't need the result of $\frac{\partial f}{\partial y}$.

[jakevdp]

I may not be understanding your question, but given a function with two inputs and one output, grad can be used to return the partial derivative with respect to only one output:

import jax

def f(x, y):
  return x * y

df_dx = jax.grad(f, argnums=0)(1.0, 2.0)
print(df_dx) # 2.0

I can't think of any more efficient way to do this with JAX.

[wuwudawen]

I've realized that it's possible to isolate a single scalar component from a function's input vector.
Consider the function:

def f(x):
    return jnp.square(x).sum()

Here, $x=[x_1, x_2, \dots, x_n]\in\mathbb{R}^n$.
I'm interested in computing $\frac{\partial f}{\partial x_i}$ for a particular $x_i$.

To achieve this, redefine the function as:

def f(xpi, xi):
    return jnp.square(xpi).sum() + jnp.square(xi)
result = grad(f, argnums=1)(xpi, xi)

Here $xpi\in\mathbb{R}^{n-1}$ represents the vector $x$ with the i-th component removed, and $x_i\in\mathbb{R}$ is the i-th component. This approach works but requires the function to be restructured.

I wonder if there's a more convenient method, perhaps one that allows specifying the component of interest directly when computing the gradient, like using an argument structure similar to argnums=(None, 3) to indicate differentiation with respect to the 3rd component of the second input vector.

[carlosgmartin]

This is something I'd be interested in as well (the ability to specify a subtree to take the gradient with respect to in jax.grad).

[jakevdp]

I think the question here is somewhat different: not how to take the gradient with respect to a subtree, but rather taking the gradient with respect to a single element of an array.

If you're looking for how to take gradients with respect to an element or elements of a pytree, I'd start with #20288 and discussions referenced therein.