Repository: google/jax

Issue:
eliminate use of `jnp.ravel` in `jnp.histogram`
In general, `jnp.ravel` tends to behave extremely badly in a distributed context, and should probably be avoided wherever possible - it will at least create large amounts of ICI traffic, and at worst produce giant replicated tensors on every chip.

I've got a draft implementation that isn't quite equivalent enough for a PR.

```py
def bincount(
    x,
    *,
    length: int,
    weights=None,
) -> Array:
    if not jnp.issubdtype(x.dtype, jnp.integer):
        raise TypeError(
            f"x argument to bincount must have an integer type; got {x.dtype}"
        )
    if weights is None:
        weights = np.array(1, dtype=jnp.int_)
    return jnp.zeros(length, weights.dtype).at[jnp.clip(x, 0)].add(weights)


def efficient_hist(
    x: jax.Array,
    bins: int = 100,
    density: bool = False,
):
    min = jnp.min(x)
    max = jnp.max(x)
    bin_edges = jnp.linspace(min, max, bins + 1)
    bin_idx = jnp.searchsorted(bin_edges, x, side="right")
    bin_idx = jnp.clip(bin_idx, 0, len(bin_edges) - 1)
    counts = bincount(bin_idx, length=len(bin_edges))[1:]
    if density:
        bin_widths = jnp.diff(bin_edges)
        counts = counts / bin_widths / jnp.sum(counts)
    return counts, bin_edges
```

I needed the new `bincount` function because the current implementation has an assertion that `x` is one dimensional which is entirely unnecessary.

I've tested and this implementation works without issue in a distributed setting where the intermediate tensors are too large to construct on any single chip.


Discussion:
Thanks - the reason that `jax.numpy.bincount` raises an error for two-dimensional inputs is that `numpy.bincount` also raises an error in such cases, and the semantics of `jax.numpy` functions by design match the semantics of `numpy` functions.

That said, if this more efficient histogram implementation passes existing tests, we'd be happy to change the implementation. Are you interested in putting together a pull request?
Since `efficient_hist` accepts only evenly spaced bins, you can replace `searchsorted` with `bin_idx = (x - min) // ((max - min) / bins)`

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -423,7 +423,7 @@ def histogram_bin_edges(a: ArrayLike, bins: ArrayLike = 10,
   if isinstance(bins, str):
     raise NotImplementedError("string values for `bins` not implemented.")
   util.check_arraylike("histogram_bin_edges", a, bins)
-  arr = ravel(a)
+  arr = asarray(a)
   dtype = dtypes.to_inexact_dtype(arr.dtype)
   if _ndim(bins) == 1:
     return asarray(bins, dtype=dtype)
@@ -448,18 +448,18 @@ def histogram(a: ArrayLike, bins: ArrayLike = 10,
               density: bool | None = None) -> tuple[Array, Array]:
   if weights is None:
     util.check_arraylike("histogram", a, bins)
-    a = ravel(*util.promote_dtypes_inexact(a))
+    a, = util.promote_dtypes_inexact(a)
     weights = ones_like(a)
   else:
     util.check_arraylike("histogram", a, bins, weights)
     if shape(a) != shape(weights):
       raise ValueError("weights should have the same shape as a.")
-    a, weights = map(ravel, util.promote_dtypes_inexact(a, weights))
+    a, weights = util.promote_dtypes_inexact(a, weights)
 
   bin_edges = histogram_bin_edges(a, bins, range, weights)
   bin_idx = searchsorted(bin_edges, a, side='right')
   bin_idx = where(a == bin_edges[-1], len(bin_edges) - 1, bin_idx)
-  counts = bincount(bin_idx, weights, length=len(bin_edges))[1:]
+  counts = zeros(len(bin_edges), weights.dtype).at[bin_idx].add(weights)[1:]
   if density:
     bin_widths = diff(bin_edges)
     counts = counts / bin_widths / counts.sum()


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2808,7 +2808,7 @@ def testHistogramBinEdges(self, shape, dtype, bins, range, weights):
                           atol=tol, rtol=tol)
 
   @jtu.sample_product(
-    shape=[(5,), (5, 5)],
+    shape=[(5,), (4, 5)],
     dtype=default_dtypes,
     # We only test explicit integer-valued bin edges because in other cases
     # rounding errors lead to flaky tests.
@@ -2819,17 +2819,17 @@ def testHistogramBinEdges(self, shape, dtype, bins, range, weights):
   def testHistogram(self, shape, dtype, bins, density, weights):
     rng = jtu.rand_default(self.rng())
     _weights = lambda w: abs(w) if weights else None
-    np_fun = lambda a, w: np.histogram(a, bins=bins, density=density,
-                                         weights=_weights(w))
+    def np_fun(a, w):
+      # Numpy can't handle bfloat16
+      a = a.astype('float32') if a.dtype == jnp.bfloat16 else a
+      w = w.astype('float32') if w.dtype == jnp.bfloat16 else w
+      return np.histogram(a, bins=bins, density=density, weights=_weights(w))
     jnp_fun = lambda a, w: jnp.histogram(a, bins=bins, density=density,
                                          weights=_weights(w))
     args_maker = lambda: [rng(shape, dtype), rng(shape, dtype)]
     tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
-    # np.searchsorted errors on bfloat16 with
-    # "TypeError: invalid type promotion with custom data type"
-    if dtype != jnp.bfloat16:
-      self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
-                              tol=tol)
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False,
+                            tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @jtu.sample_product(