autograd

Autograd v1.2 introduced a new interface for defining custom vector-Jacobian products (VJPs). If you write custom primitives, you must migrate from using the .defvjp() method on a primitive object to using the defvjp function from autograd.extend.

Key changes include:

• defvjp is now a standalone function: defvjp(func, ...) instead of func.defvjp(...).
• VJPs are now 'staged': instead of a single lambda accepting the gradient g, you provide a function that returns a lambda: lambda ans, *args: lambda g: ....
• The vs (values) and gvs (gradients) arguments have been removed to improve performance and memory efficiency.
• argnums are handled via a specific parameter in the defvjp function call.

import autograd.numpy as np
from autograd.extend import primitive, defvjp

@primitive
def func(x, y, z):
    assert z != 0
    return x * y**2

# New way to define VJPs
defvjp(func,
       lambda ans, x, y, z: lambda g: g * y**2,
       lambda ans, x, y, z: lambda g: 2 * g * x * y,
       None)

Autograd supports most mathematical operations, array/matrix manipulation, N-dimensional convolutions, and FFT routines. It provides full support for complex numbers and some scipy.stats.norm routines.

Best Practices

Do use:

• Most autograd.numpy functions.
• Most numpy.ndarray methods.
• Indexing and slicing (e.g., x = A[3, :, 2:4]).
• Explicit array creation from lists (e.g., A = np.array([x, y])).

Don't use:

• Assignment to arrays: A[0,0] = x is not supported for arrays being differentiated.
• Implicit list casting: Avoid A = np.sum([x, y]). Instead, use A = np.sum(np.array([x, y])).
• Method notation for dot products: Avoid A.dot(B). Use np.dot(A, B) instead.
• In-place operations: Avoid a += b. Use a = a + b instead.
• Standard isinstance checks: Standard isinstance(x, np.ndarray) may return False. Use from autograd.builtins import isinstance instead.

Autograd works with any container implementing the NumPy __array_ufunc__ protocol (e.g., xarray.DataArray). You can use standard autograd.numpy functions directly on these containers.

Important: Since grad requires a scalar output, you must extract the underlying array (e.g., using .data for Xarray) before performing the final reduction in your loss function.

import autograd.numpy as np
from autograd import grad
import xarray as xr

# A named-axis dataset
measurements = xr.DataArray(
    np.array([[0.5, 1.2, -0.3],
              [0.1, 0.4,  0.9],
              [-0.7, 0.2, 1.1],
              [0.3, -0.5, 0.6]]),
    dims=["time", "feature"],
    coords={"feature": ["a", "b", "c"]},
)

def loss(weights):
    # Broadcasting respects the named "feature" axis
    scores = np.tanh(measurements * weights)
    # Pull the plain array back out for the scalar reduction grad needs
    return np.sum(scores.data ** 2)

weights = np.array([0.5, -1.0, 2.0])
print("loss:", loss(weights))
print("grad:", grad(loss)(weights))

If Autograd does not support a specific function (e.g., an external C call or a custom stable math function), you can extend it by defining a new primitive.

1. Define the function: Use the @primitive decorator. This tells Autograd to treat the function as a black box.
2. Define the VJP: Write a function that returns a vector-Jacobian product (VJP) operator. This operator should take the output gradient g and return a function that computes the gradient with respect to the inputs.
3. Register the primitive: Use defvjp(function, vjp_function) to link them.

To support higher-order derivatives, ensure the code inside the VJP function is itself differentiable by Autograd (i.e., written using Autograd primitives).

import autograd.numpy as np
from autograd.extend import primitive, defvjp
from autograd import grad

# 1. Define the function
@primitive
def logsumexp(x):
    """Numerically stable log(sum(exp(x)))"""
    max_x = np.max(x)
    return max_x + np.log(np.sum(np.exp(x - max_x)))

# 2. Define the VJP
def logsumexp_vjp(ans, x):
    x_shape = x.shape
    return lambda g: np.full(x_shape, g) * np.exp(x - np.full(x_shape, ans))

# 3. Register
defvjp(logsumexp, logsumexp_vjp)

# Usage
def example_func(y):
    z = y**2
    lse = logsumexp(z)
    return np.sum(lse)

grad_of_example = grad(example_func)
print("Gradient: ", grad_of_example(np.array([1.5, 6.7, 1e-10])))

This example demonstrates a complete workflow: defining a model (sigmoid/logistic), a loss function (negative log-likelihood), and optimizing weights using gradient descent.

import autograd.numpy as np
from autograd import grad

def sigmoid(x):
    return 0.5 * (np.tanh(x / 2.) + 1)

def logistic_predictions(weights, inputs):
    return sigmoid(np.dot(inputs, weights))

def training_loss(weights):
    preds = logistic_predictions(weights, inputs)
    label_probabilities = preds * targets + (1 - preds) * (1 - targets)
    return -np.sum(np.log(label_probabilities))

# Build a toy dataset.
inputs = np.array([[0.52, 1.12,  0.77],
                   [0.88, -1.08, 0.15],
                   [0.52, 0.06, -1.30],
                   [0.74, -2.49, 1.39]])
targets = np.array([True, True, False, True])

# Define a function that returns gradients of training loss using Autograd.
training_gradient_fun = grad(training_loss)

# Optimize weights using gradient descent.
weights = np.array([0.0, 0.0, 0.0])
print("Initial loss:", training_loss(weights))
for i in range(100):
    weights -= training_gradient_fun(weights) * 0.01

print("Trained loss:", training_loss(weights))

The grad function is the primary tool for computing derivatives. It takes a function as input and returns a new function that computes the gradient of the input function.

Requirements:

• The input function must have a scalar-valued output (e.g., a float).
• You should use autograd.numpy (a thinly-wrapped version of Numpy) within your functions to ensure compatibility.

Autograd supports standard Python control flow like while loops, if statements, and recursion because it tracks the specific operations applied to inputs during each function call.

import autograd.numpy as np
from autograd import grad

def taylor_sine(x):
    ans = currterm = x
    i = 0
    while np.abs(currterm) > 0.001:
        currterm = -currterm * x**2 / ((2 * i + 3) * (2 * i + 2))
        ans = ans + currterm
        i += 1
    return ans

grad_sine = grad(taylor_sine)
print("Gradient of sin(pi) is", grad_sine(np.pi))