Differentiating with Respect to Arguments

By default, jax.grad computes the gradient of a function with respect to its first argument. However, many functions, especially in machine learning contexts like loss functions, take multiple inputs (e.g., parameters and data). You often need to calculate the gradient with respect to only one or a specific subset of these inputs. JAX provides the argnums parameter within jax.grad to control this behavior precisely.

Specifying a Single Argument for Differentiation

Let's consider a simple function that takes two arguments, $f(x, y) = x^2 \cdot y$.

import jax
import jax.numpy as jnp

def power_product(x, y):
  """Computes x^2 * y"""
  return (x**2) * y

# Define some input values
x_val = 3.0
y_val = 4.0

print(f"Function output: {power_product(x_val, y_val)}")

If we apply jax.grad without any extra parameters, it differentiates with respect to the first argument, $x$. The partial derivative $\frac{\partial f}{\partial x}$ is $2xy$.

# Gradient with respect to the first argument (x) - default behavior
grad_f_wrt_x = jax.grad(power_product)
gradient_x = grad_f_wrt_x(x_val, y_val)

print(f"Gradient w.r.t. x (default): {gradient_x}") # Expected: 2 * 3.0 * 4.0 = 24.0

Now, suppose we need the gradient with respect to the second argument, $y$. The partial derivative $\frac{\partial f}{\partial y}$ is $x^2$. We can achieve this using argnums=1 (remembering that argument indexing is 0-based):

# Gradient with respect to the second argument (y) using argnums=1
grad_f_wrt_y = jax.grad(power_product, argnums=1)
gradient_y = grad_f_wrt_y(x_val, y_val)

print(f"Gradient w.r.t. y (argnums=1): {gradient_y}") # Expected: 3.0**2 = 9.0

Here, argnums=1 tells jax.grad to compute the derivative considering the second argument (y in this case) as the variable and treating other arguments (x) as constants for the differentiation process.

Specifying Multiple Arguments for Differentiation

You might need gradients with respect to several arguments simultaneously. For instance, consider a function $g(w, b, x) = wx + b$. We might want the gradients with respect to both $w$ and $b$.

You can achieve this by passing a tuple of integers to argnums.

def affine(w, b, x):
  """Computes w*x + b"""
  return w * x + b

# Define input values
w_val = 2.0
b_val = 1.0
x_data = 5.0

print(f"Affine function output: {affine(w_val, b_val, x_data)}")

# Gradient with respect to the first (w) and second (b) arguments
grad_g_wrt_wb = jax.grad(affine, argnums=(0, 1))

# Note: The input x_data is treated as constant during differentiation
gradients_wb = grad_g_wrt_wb(w_val, b_val, x_data)

print(f"Gradient w.r.t. (w, b) using argnums=(0, 1): {gradients_wb}")
# Expected: (gradient w.r.t w, gradient w.r.t b) = (x, 1) = (5.0, 1.0)

When argnums is a tuple, the function returned by jax.grad also returns a tuple. The elements of the output tuple correspond directly to the arguments specified in argnums, in the same order. In the example above:

• The first element of gradients_wb (which is 5.0) is $\frac{\partial g}{\partial w} = x$.
• The second element (which is 1.0) is $\frac{\partial g}{\partial b} = 1$.

Connection to Machine Learning

This capability is fundamental in training machine learning models. A typical loss function might look like loss(params, data_batch). To update the model parameters using gradient descent, you need the gradient of the loss with respect to params, while treating data_batch as fixed inputs. This is naturally expressed as:

# Example structure
# def loss_function(params, data_batch):
#   predictions = model_apply(params, data_batch['inputs'])
#   error = predictions - data_batch['targets']
#   return jnp.mean(error**2) # Example: Mean Squared Error

# grad_loss_wrt_params = jax.grad(loss_function, argnums=0)
# gradients = grad_loss_wrt_params(current_params, batch)
# updated_params = current_params - learning_rate * gradients

Using argnums=0 ensures that jax.grad computes $\nabla_{\text{params}} \text{loss}(\text{params}, \text{data\_batch})$, which is exactly what's needed for optimization.

By mastering the argnums parameter, you gain fine-grained control over the automatic differentiation process in JAX, allowing you to target specific inputs for gradient computation, which is essential for more complex functions and standard machine learning workflows.