scipy.differentiate._differentiate:jacobian

Notes

Suppose we wish to evaluate the Jacobian of a function $f:{\mathbf{R}}^m\to {\mathbf{R}}^n$. Assign to variables m and n the positive integer values of $m$ and $n$, respectively, and let ... represent an arbitrary tuple of integers. If we wish to evaluate the Jacobian at a single point, then:

• argument x must be an array of shape (m,)

• argument f must be vectorized to accept an array of shape (m, ...). The first axis represents the $m$ inputs of $f$; the remainder are for evaluating the function at multiple points in a single call.

• argument f must return an array of shape (n, ...). The first axis represents the $n$ outputs of $f$; the remainder are for the result of evaluating the function at multiple points.

• attribute df of the result object will be an array of shape (n, m), the Jacobian.

This function is also vectorized in the sense that the Jacobian can be evaluated at k points in a single call. In this case, x would be an array of shape (m, k), f would accept an array of shape (m, k, ...) and return an array of shape (n, k, ...), and the df attribute of the result would have shape (n, m, k).

Suppose the desired callable f_not_vectorized is not vectorized; it can only accept an array of shape (m,). A simple solution to satisfy the required interface is to wrap f_not_vectorized as follows

def f(x):
    return np.apply_along_axis(f_not_vectorized, axis=0, arr=x)

Alternatively, suppose the desired callable f_vec_q is vectorized, but only for 2-D arrays of shape (m, q). To satisfy the required interface, consider

def f(x):
    m, batch = x.shape[0], x.shape[1:]  # x.shape is (m, ...)
    x = np.reshape(x, (m, -1))  # `-1` is short for q = prod(batch)
    res = f_vec_q(x)  # pass shape (m, q) to function
    n = res.shape[0]
    return np.reshape(res, (n,) + batch)  # return shape (n, ...)

Then pass the wrapped callable f as the first argument of jacobian.

Array API Standard Support

jacobian has experimental support for Python Array API Standard compatible backends in addition to NumPy. Please consider testing these features by setting an environment variable SCIPY_ARRAY_API=1 and providing CuPy, PyTorch, JAX, or Dask arrays as array arguments. The following combinations of backend and device (or other capability) are supported.

====================  ====================  ====================
Library               CPU                   GPU
====================  ====================  ====================
NumPy                 ✅                     n/a                 
CuPy                  n/a                   ✅                   
PyTorch               ✅                     ✅                   
JAX                   ⚠️ no JIT             ⚠️ no JIT           
Dask                  ⛔                     n/a                 
====================  ====================  ====================

See dev-arrayapi for more information.