jax-ml/jax · error · NotImplementedError

only 1D box integrations are supported; use `integrate_box_1

Error message

only 1D box integrations are supported; use `integrate_box_1d`

What it means

Unlike scipy, JAX's gaussian_kde does not implement the general N-D integrate_box; it is stubbed to raise NotImplementedError and direct users to integrate_box_1d. This is an intentional API-parity gap documented in the method docstring.

Source

Thrown at jax/_src/scipy/stats/kde.py:223

                                     dtype=self.dataset.dtype).T
    return self.dataset[:, ind] + eps

  def pdf(self, x):
    """Probability density function"""
    return self.evaluate(x)

  def logpdf(self, x):
    """Log probability density function"""
    check_arraylike("logpdf", x)
    x = self._reshape_points(x)
    result = _gaussian_kernel_eval(True, self.dataset.T, self.weights[:, None],
                                   x.T, self.inv_cov)
    return result[:, 0]

  def integrate_box(self, low_bounds, high_bounds, maxpts=None):
    """This method is not implemented in the JAX interface."""
    del low_bounds, high_bounds, maxpts
    raise NotImplementedError(
        "only 1D box integrations are supported; use `integrate_box_1d`")

  def set_bandwidth(self, bw_method=None):
    """This method is not implemented in the JAX interface."""
    del bw_method
    raise NotImplementedError(
        "dynamically changing the bandwidth method is not supported")

  def _reshape_points(self, points):
    if dtypes.issubdtype(lax.dtype(points), np.complexfloating):
      raise NotImplementedError(
          "gaussian_kde does not support complex coordinates")
    points = jnp.atleast_2d(points)
    d, m = points.shape
    if d != self.d:
      if d == 1 and m == self.d:
        points = jnp.reshape(points, (self.d, 1))
      else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

For 1-D KDEs use kde.integrate_box_1d(low, high)
For N-D, estimate via Monte Carlo: sample with kde.resample and count points in the box
Alternatively evaluate kde on a grid and numerically integrate (trapezoid/quad)

Example fix

// before
p = kde.integrate_box(lo, hi)
// after (1-D)
p = kde.integrate_box_1d(lo, hi)

Defensive patterns

Strategy: fallback

Validate before calling

hasattr check not needed; branch on dimensionality:
use_box_1d = kde.d == 1

Try / catch

try:
    p = kde.integrate_box(lo, hi)
except NotImplementedError:
    samples = kde.resample(200_000, seed=key)
    p = ((samples >= lo[:, None]) & (samples <= hi[:, None])).all(0).mean()

Prevention

Replace integrate_box with integrate_box_1d when d == 1
Keep a Monte Carlo helper for N-D boxes
Check JAX docs for scipy API gaps before porting

When it happens

Trigger: Calling kde.integrate_box(low, high) for any KDE, even 1-D ones.

Common situations: Porting scipy.stats.gaussian_kde code that used integrate_box for multidimensional probability mass; assuming jax.scipy mirrors the full scipy API.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/729da3fe134a137b. Report an issue: GitHub.