liveBook · Manning

Argnums

As explained in Deep Learning with JAX

The grad-related parameter that selects which positional argument or arguments the derivative should be taken with respect to.

argnums chooses which positional argument or arguments a JAX gradient should differentiate with respect to. By default, grad() uses the first parameter, so argnums matters when the differentiated input is not first or when you want several inputs at once. section 4.2.3 section 4.4

Common uses here include:

  • Skipping a non-differentiated leading parameter without rewriting the function signature. section 4.2.3
  • Differentiating with respect to multiple positional inputs by passing a tuple of indices. section 4.2.3
  • Working with structured inputs such as dicts and pytrees through the selected argument. section 4.2.3

Selecting the second argument

For a distance function written as dist(order, x, y), argnums=1 targets x without moving order to the end. The result is a gradient shaped like x, while argnums=(1,2) returns a tuple of gradients for both x and y.

argnums works alongside other gradient options: has_aux returns extra outputs, value_and_grad returns both loss and gradient, and stop_gradient blocks flow through selected expressions.

How Deep Learning with JAX covers it

How deep it goes
Explains integer and tuple forms, return shapes, and a few related gradient options, but stops short of internals.
Examples
Uses a distance function, a linear-model training step, and a gradient-blocking branch example.
What it assumes
Assumes basic grad() usage and positional function arguments, while teaching the selection behavior from zero.
What you can do after
Apply it via grad() to target one or several inputs and receive matching gradient shapes.
How it connects
Links grad(), value_and_grad(), has_aux, pytree, dicts, and stop_gradient.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.