jax-ml/jax · error · ValueError

the limits of integration in integrate_box_1d must be scalar

Error message

the limits of integration in integrate_box_1d must be scalars

What it means

integrate_box_1d requires low and high to be scalar (np.ndim == 0). Passing arrays, lists, or length-1 arrays as integration limits raises this error, since the method computes elementwise (low - dataset)/sigma and needs a single pair of bounds.

Source

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

    if mean.shape != (self.d,):
      raise ValueError(f"mean does not have dimension {self.d}")
    if cov.shape != (self.d, self.d):
      raise ValueError(f"covariance does not have dimension {self.d}")

    chol = linalg.cho_factor(self.covariance + cov)
    norm = jnp.sqrt(2 * np.pi)**self.d * jnp.prod(jnp.diag(chol[0]))
    norm = 1.0 / norm
    return _gaussian_kernel_convolve(chol, norm, self.dataset, self.weights,
                                     mean)

  @api.jit
  def integrate_box_1d(self, low, high):
    """Integrate the distribution over the given limits."""
    if self.d != 1:
      raise ValueError("integrate_box_1d() only handles 1D pdfs")
    if np.ndim(low) != 0 or np.ndim(high) != 0:
      raise ValueError(
          "the limits of integration in integrate_box_1d must be scalars")
    sigma = jnp.squeeze(jnp.sqrt(self.covariance))
    low = jnp.squeeze((low - self.dataset) / sigma)
    high = jnp.squeeze((high - self.dataset) / sigma)
    return jnp.sum(self.weights * (special.ndtr(high) - special.ndtr(low)))

  def integrate_kde(self, other):
    """Integrate the product of two Gaussian KDE distributions."""
    if other.d != self.d:
      raise ValueError("KDEs are not the same dimensionality")

    chol = linalg.cho_factor(self.covariance + other.covariance)
    norm = jnp.sqrt(2 * np.pi)**self.d * jnp.prod(jnp.diag(chol[0]))
    norm = 1.0 / norm

    sm, lg = (self, other) if self.n < other.n else (other, self)
    result = api.vmap(partial(_gaussian_kernel_convolve, chol, norm, lg.dataset,
                              lg.weights),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Extract scalars: float(low), float(high), or low.item() if tracing is not required
  2. Use jax.vmap over scalar-bound calls for many intervals
  3. Squeeze shape-(1,) bounds with jnp.squeeze before calling

Example fix

// before
p = kde.integrate_box_1d(bounds[0], bounds[1])  # bounds[i].shape == (1,)
// after
p = kde.integrate_box_1d(jnp.squeeze(bounds[0]), jnp.squeeze(bounds[1]))
Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
assert np.ndim(low) == 0 and np.ndim(high) == 0

Type guard

def scalar_bounds(low, high) -> bool:
    return np.ndim(low) == 0 and np.ndim(high) == 0

Prevention

When it happens

Trigger: kde.integrate_box_1d(jnp.array([0.0]), 1.0), or passing arrays of bounds hoping for batched integrals.

Common situations: Limits coming from a config array or a previous computation that returns shape (1,) tensors; trying to vectorize box integrals over many interval pairs.

Related errors


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