Gaps in NumPy API coverage: unwrapped subpackages and missing gradients

I asked Claude Code to help prepare this document after some prompting about feature gaps, and manually validated its output below. Here is the AI output reproduced below

Summary

Autograd says it wraps the whole NumPy API. That is true if you only mean "can you call the function", because wrap_namespace turns every public function in numpy, numpy.fft, numpy.linalg, and numpy.random into a primitive. I checked against NumPy 2.4.6 and there are zero missing functions across all four namespaces.

The real gaps show up once you try to differentiate. There are two kinds:

1. Subpackages that never get wrapped. Today only fft, linalg, and random are re-exported as autograd-aware versions. A few other NumPy subpackages that hold differentiable functions are simply not there.
2. Functions that are wrapped but have no gradient. They run fine in the forward pass and then raise NotImplementedError: VJP of <fn> wrt argnums ... not defined as soon as you call grad through them. Around 50 of these are functions that really ought to be differentiable.

This issue tracks both. I have left out everything that does not need a gradient in the first place, such as constructors like zeros and eye, I/O like save and loadtxt, integer and bitwise ops, comparisons, boolean predicates like isnan, the sort and search functions that return indices like argmax and argsort, set operations, and rounding like floor and round.

Why are these missing?

Most of autograd's gradient work happened between 2017 and 2019. Of the 59 commits that ever touched numpy_vjps.py, 48 land in that window. After that the project went quiet on new features as attention shifted to JAX, and the releases since then have mostly been about packaging and keeping up with newer Python and NumPy versions rather than adding gradients.

That history explains the two flavours of gap. The first flavour is everything that NumPy only added in the 2.x series, from 2024 onward. None of it existed while autograd was being actively built, so there was nothing to write a gradient for. This covers the array-API style functions (cumulative_sum, cumulative_prod, vecdot, matvec, vecmat, matrix_transpose, unstack), the rename of trapz to trapezoid, and the new numpy.strings namespace. Worth noting: NumPy usually shipped these as brand new function objects rather than as aliases, so even where autograd already knows how to differentiate the underlying operation (cumsum, for instance), that gradient does not automatically apply to the new name.

The second flavour is the functions that were already around during active development but never got picked up. The NaN-aware reductions, average, median, interp, convolve, and the polynomial, ma, and emath subpackages all fall here. They were skipped because the effort went into the operations that show up most often in optimisation and deep learning code, not because they were hard.

Here are the introduction versions, all confirmed against NumPy's own release notes, so the old versus new split is easy to see.

Where to start: the best value for the least effort

If we do this in stages, here is the order I would suggest.

Stage 1, the quick wins. These are the NumPy 2.x array-API names whose underlying operation already has a gradient in autograd. They mostly need to be pointed at the gradient that already exists, or rewritten as a thin wrapper over the operation that already works. There is almost no new maths involved, and it unblocks anyone writing array-API style code on NumPy 2.x. The list is cumulative_sum, cumulative_prod, matrix_transpose, vecdot, and unstack. matvec and vecmat are a little more involved but still build straight on top of matmul and dot.

Stage 2, the NaN reductions. This is a whole family that people hit all the time and that currently fails as a block. Each one is a small variation on a reduction whose gradient autograd already knows, such as sum, prod, mean, max, min, cumsum, and so on. The gradient is the same as the non-NaN version, except the NaN positions get zero gradient. Implementing the first one sets the pattern for all the rest. The set is nansum, nanprod, nanmean, nanmax, nanmin, nanstd, nanvar, nancumsum, and nancumprod, plus the selection-based nanmedian, nanpercentile, and nanquantile, which follow the median, percentile, and quantile pattern in the next stage.

Stage 3, common statistics and a few elementwise functions. These come up constantly in statistics and optimisation code. average is a weighted mean and is straightforward. median, percentile, and quantile route the gradient through whichever element gets selected. Then there is cov, corrcoef, ptp, and the elementwise cbrt, float_power, and positive.

Stage 4, useful but more work or narrower in audience. interp (piecewise linear), convolve and correlate (linear, but with axis and mode bookkeeping to get right), trapezoid, unwrap, sinc, i0, and the gather-style take, take_along_axis, compress, extract, and choose.

The subpackages are the biggest jobs. See the next section. The short version is that they are not worth doing.

When checked against JAX?

JAX is written by the same people who wrote autograd, and it is their production-grade autodiff library, so its choices are the strongest signal we have about what is worth supporting. I installed JAX 0.10.2 and tested every item on this list, both whether the function exists in jax.numpy and whether you can actually take a gradient through it. Two clear results came out of that.

First, every single function-level gap below exists in jax.numpy and returns a finite gradient. That confirms all of them are both feasible and worth doing. The only question is the order.

Second, and this is the important one, JAX ships none of the missing subpackages. There is no jax.numpy.ma, no jax.numpy.polynomial, no jax.numpy.emath, and no jax.numpy.strings. The same team that built autograd had ten years and every incentive to port these and chose not to.

• numpy.ma (masked arrays). Do not add. Masking is stateful, shape-dynamic, and built around in-place mask manipulation, none of which fits a functional, trace-based autodiff model. The idiomatic way to ignore entries in an autodiff setting is already where plus the NaN-aware reductions, or a multiplicative mask.
• numpy.strings. Not relevant. String operations have no derivative. It only showed up here because it is a subpackage. Drop it.
• numpy.polynomial. Very low priority. It is a large class-based surface, and the genuinely useful part (fitting and evaluating polynomials) already lives in autograd's main namespace as polyfit, polyval, and friends. JAX never built the class package.
• numpy.emath. Very low priority. Its whole job is to switch into the complex domain for out-of-range real inputs, such as sqrt(-1) returning 1j. Differentiating across that branch switch is awkward and the use case is niche.

On the function side, JAX makes all of them differentiable, but the value to a real user is not uniform. The checklist below is grouped to reflect that. A few notes worth keeping in mind:

• The selection-based functions (median, quantile, percentile, and their nan variants, plus nanmax and nanmin) return a subgradient that routes to the selected element or the interpolated pair. This is mathematically valid and exactly what JAX returns, the same way max-pooling has a subgradient, but it is not smooth. Worth a line in the docs when these land.
• sinc needs a custom rule at x = 0 to avoid a 0/0. JAX handles this with a dedicated JVP, and autograd should do the same.
• A handful of functions are cheap to add but rarely matter for differentiation, so they should not be prioritised: positive is the identity, float_power is a dtype-promoting duplicate of power, ptp is just max - min, and choose, compress, extract, and resize use dynamic or boolean-mask semantics that seldom appear inside a differentiable objective. They are marked below.

Checklist

(A) Subpackages that are not wrapped at all (recommend not adding)

These are missing from autograd.numpy, but as explained above, JAX omits all of them and so should we. Listed here for completeness and to record the decision, not as work to pick up.

• numpy.ma. Masked arrays. Recommend not adding. Stateful, shape-dynamic masking does not fit trace-based autodiff. Use where plus the NaN reductions instead.
• numpy.strings. Not relevant, no derivatives exist. Recommend dropping from scope.
• numpy.polynomial. Recommend not adding. Large class-based surface, and the useful fitting and evaluation functions already exist in the main namespace.
• numpy.emath. Recommend not adding. Niche complex-domain branch switching.

(B) Stage 1: array-API names whose underlying op already has a gradient

• cumulative_sum (mirror cumsum)
• cumulative_prod (mirror cumprod)
• matrix_transpose
• vecdot
• unstack
• matvec
• vecmat

(C) Stage 2: NaN-aware reductions

• nansum
• nanprod
• nanmean
• nanmax
• nanmin
• nanstd
• nanvar
• nancumsum
• nancumprod
• nanmedian
• nanpercentile
• nanquantile

(D) Stage 3: common statistics and elementwise

• average (weighted mean)
• median
• percentile
• quantile
• cov
• corrcoef
• cbrt
• ptp (low value, just max - min)
• float_power (low value, behaves like power)
• positive (low value, identity function)

(E) Stage 4: signal, interpolation, and gather

• interp (https://github.com/HIPS/autograd/issues/193)
• convolve
• correlate
• trapezoid
• unwrap
• sinc (needs a custom rule at x = 0)
• i0
• take
• take_along_axis
• compress (low value, boolean-mask gather)
• extract (low value, boolean-mask gather)
• choose (low value, rarely differentiated)

(F) Constructive and linear helpers, lower priority but still differentiable

• diagflat
• vander
• kron
• block
• dstack
• meshgrid (https://github.com/HIPS/autograd/issues/457)
• resize (low value, tiling/truncation gradient rarely needed)

How to reproduce

Forward coverage, which should report nothing missing in any namespace:

import numpy as np
import autograd.numpy as anp

for mod in [None, "fft", "linalg", "random"]:
    npmod = np if mod is None else getattr(np, mod)
    amod = anp if mod is None else getattr(anp, mod)
    pub = [n for n in npmod.__all__ if not n.startswith("_")]
    missing = [n for n in pub if not hasattr(amod, n)]
    print(mod or "numpy", "missing:", missing)

A function that is wrapped but cannot be differentiated, which should raise:

from autograd import grad
grad(lambda x: anp.nansum(x))(np.array([1.0, 2.0, np.nan]))
# NotImplementedError: VJP of nansum wrt argnums (0,) not defined

To check what JAX does with the same functions, without installing anything permanently:

uvx --with jax python -c "
import jax.numpy as jnp
from jax import grad
import numpy as np
x = jnp.array([1.0, 2.0, float('nan')])
...
"

A couple of notes:

• Checked against NumPy 2.4.6 and JAX 0.10.2.
• Deliberately left out of the lists above as not needing a gradient: constructors, I/O, integer and bitwise ops, comparisons, boolean predicates, the sort and search functions that return indices, set operations, and rounding.