argnums, value_and_grad, and has_aux

The basic jax.grad(f) differentiates with respect to the first argument. But real-world functions often need more control.

argnums: Choosing What to Differentiate

import jax
import jax.numpy as jnp

def f(x, y, z):
    return x ** 2 + y ** 3 + z

# Differentiate with respect to x (argument 0) — default
df_dx = jax.grad(f, argnums=0)
print(df_dx(2.0, 3.0, 4.0))  # 4.0 (d/dx of x^2 = 2x = 4)

# Differentiate with respect to y (argument 1)
df_dy = jax.grad(f, argnums=1)
print(df_dy(2.0, 3.0, 4.0))  # 27.0 (d/dy of y^3 = 3y^2 = 27)

# Differentiate with respect to multiple arguments
df_dxy = jax.grad(f, argnums=(0, 1))
grads = df_dxy(2.0, 3.0, 4.0)
print(grads)  # (4.0, 27.0) — tuple of gradients

value_and_grad: Get Both in One Pass

In training loops, you need both the loss value (for logging) and its gradient (for updates). Computing them separately wastes work. jax.value_and_grad computes both in a single forward+backward pass:

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

# Returns (value, gradient) tuple
val_and_grad_fn = jax.value_and_grad(loss_fn)

params = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [0.3, 0.8]])
y = jnp.array([1.0, 0.5])

loss_val, grads = val_and_grad_fn(params, x, y)
print(f"Loss: {loss_val:.4f}")
print(f"Gradients: {grads}")

has_aux: Returning Extra Data Alongside Gradients

Sometimes your loss function needs to return additional information (metrics, intermediate activations, etc.) alongside the scalar loss. Use has_aux=True to tell JAX that the function returns (loss, auxiliary_data):

import jax
import jax.numpy as jnp

def loss_with_metrics(params, x, y):
    pred = jnp.dot(x, params)
    loss = jnp.mean((pred - y) ** 2)
    # Return loss AND extra info
    metrics = {
        'mse': loss,
        'predictions': pred,
        'max_error': jnp.max(jnp.abs(pred - y))
    }
    return loss, metrics  # (scalar, auxiliary)

# has_aux=True tells grad: the function returns (loss, aux),
# only differentiate with respect to the loss part
grad_fn = jax.grad(loss_with_metrics, has_aux=True)

params = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [0.3, 0.8]])
y = jnp.array([1.0, 0.5])

grads, metrics = grad_fn(params, x, y)
print(f"Gradients: {grads}")
print(f"MSE: {metrics['mse']:.4f}")
print(f"Max error: {metrics['max_error']:.4f}")

# Combined: value_and_grad with has_aux
val_grad_fn = jax.value_and_grad(loss_with_metrics, has_aux=True)
(loss, metrics), grads = val_grad_fn(params, x, y)