💡Add a small example to jax.grad API about taking gradients w.r.t. an n-th parameter (argnums)

[8bitmp3]

Hi @mattjj and team.

Idea

Provide a very small example that shows how to do a common task with jax.api's jax.grad like grading w.r.t. 1st and 2nd parameters (using those words) for better clarity and UX especially when a bunch of external folks are probably learning JAX and it's magic autodiff.

More details

For example, you can write:

grad(my_function, argnums=(0, 1)

The docstring already explains it like this:

def grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
         has_aux: bool = False, holomorphic: bool = False,
         allow_int: bool = False) -> Callable:
  """
    ...
  Args:
    ...
    argnums: Optional, integer or sequence of integers. Specifies which
      positional argument(s) to differentiate with respect to (default 0).
    ...

  For example:

  >>> import jax
  >>>
  >>> grad_tanh = jax.grad(jax.numpy.tanh)
  >>> print(grad_tanh(0.2))
  0.961043

  """

Searching for positional argument to differentiate in the search bar wasn't the first thing that came to my mind, although I'm not a researcher/programmer, so I may be wrong.

Proposal

Maybe we can add a second small example to the docsting explaining how to take gradient with respect to... parameter (<- key words for SEO) with some toy example like grad(my_function, argnums=(0, 1) really quick without taking up much space.

WDYT? Cheers ☕️

cc @dynamicwebpaige for UX

---

There's an example in How JAX primitives work but you'd have to search for it:

from jax import lax
from jax import api
...
def square_add_lax(a, b):
  """A square-add function using the newly defined multiply-add."""
  return multiply_add_lax(a, a, b)
...
# Differentiate w.r.t. the first argument
print("grad(square_add_lax) = ", api.grad(square_add_lax, argnums=0)(2.0, 10.))

Source: https://jax.readthedocs.io/en/latest/_modules/jax/api.html#grad