Repository: google/jax
Issue: histogram2d error when specifying `range`

```python
from jax import numpy as jnp
a = [-2,2]
b=jnp.array([a,a])
print(b)
jnp.histogram2d(jnp.ones(10), jnp.ones(10), bins=10, range=b)
```

Suggested fix from the user: [this line](https://github.com/google/jax/blob/7a40aa0114e46a7a8975e04b175dfbc7b0b388ef/jax/_src/numpy/lax_numpy.py#L1054) should probably have `range[i]`. (That seems to fix it, but I didn't try writing tests!)

@jakevdp could you take a look?

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 834e8408d..a8dd66692 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1037,6 +1037,11 @@ def histogramdd(sample, bins=10, range=None, weights=None, density=None):
   if weights is not None and weights.shape != (N,):
     raise ValueError("should have one weight for each sample.")
 
+  if range is not None and (
+      len(range) != D or _any(r is not None and len(r) != 2 for r in range)):
+    raise ValueError(f"For sample.shape={(N, D)}, range must be a sequence "
+                     f"of {D} pairs or Nones; got range={range}")
+
   try:
     num_bins = len(bins)
     if num_bins != D:
@@ -1051,7 +1056,8 @@ def histogramdd(sample, bins=10, range=None, weights=None, density=None):
   dedges = D*[None]
 
   for i in builtins.range(D):
-    bin_edges = histogram_bin_edges(sample[:, i], bins[i], range, weights)
+    range_i = None if range is None else range[i]
+    bin_edges = histogram_bin_edges(sample[:, i], bins[i], range_i, weights)
     bin_idx = searchsorted(bin_edges, sample[:, i], side='right')
     bin_idx = where(sample[:, i] == bin_edges[-1], bin_idx - 1, bin_idx)
     bin_idx_by_dim[i] = bin_idx
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 79fffd9e7..1e690b37c 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -3147,25 +3147,23 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-    {"testcase_name": "_{}_bins={}_weights={}_density={}".format(
-      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density),
-      "shape": shape,
-      "dtype": dtype,
-      "bins": bins,
-      "weights": weights,
-      "density": density
+    {"testcase_name": "_{}_bins={}_weights={}_density={}_range={}".format(
+      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density, range),
+      "shape": shape, "dtype": dtype, "bins": bins, "weights": weights, "density": density, "range": range,
     }
     for shape in [(5,), (12,)]
     for dtype in int_dtypes
     for bins in [2, [2, 2], [[0, 1, 3, 5], [0, 2, 3, 4, 6]]]
     for weights in [False, True]
     for density in [False, True]
+    for range in [None, [(-1, 1), None], [(-1, 1), (-2, 2)]]
   ))
-  def testHistogram2d(self, shape, dtype, bins, weights, density):
+  def testHistogram2d(self, shape, dtype, bins, weights, density, range):
     rng = jtu.rand_default(self.rng())
     _weights = lambda w: abs(w) if weights else None
-    np_fun = lambda a, b, w: np.histogram2d(a, b, bins=bins, weights=_weights(w), density=density)
-    jnp_fun = lambda a, b, w: jnp.histogram2d(a, b, bins=bins, weights=_weights(w), density=density)
+    np_fun = jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")(
+        lambda a, b, w: np.histogram2d(a, b, bins=bins, weights=_weights(w), density=density, range=range))
+    jnp_fun = lambda a, b, w: jnp.histogram2d(a, b, bins=bins, weights=_weights(w), density=density, range=range)
     args_maker = lambda: [rng(shape, dtype), rng(shape, dtype), rng(shape, dtype)]
     tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
     # np.searchsorted errors on bfloat16 with
@@ -3177,25 +3175,23 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-    {"testcase_name": "_{}_bins={}_weights={}_density={}".format(
-      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density),
-      "shape": shape,
-      "dtype": dtype,
-      "bins": bins,
-      "weights": weights,
-      "density": density
+    {"testcase_name": "_{}_bins={}_weights={}_density={}_range={}".format(
+      jtu.format_shape_dtype_string(shape, dtype), bins, weights, density, range),
+      "shape": shape, "dtype": dtype, "bins": bins, "weights": weights, "density": density, "range": range,
     }
     for shape in [(5, 3), (10, 3)]
     for dtype in int_dtypes
     for bins in [(2, 2, 2), [[-5, 0, 4], [-4, -1, 2], [-6, -1, 4]]]
     for weights in [False, True]
     for density in [False, True]
+    for range in [None, [(-1, 1), None, None], [(-1, 1), (-2, 2), (-3, 3)]]
   ))
-  def testHistogramdd(self, shape, dtype, bins, weights, density):
+  def testHistogramdd(self, shape, dtype, bins, weights, density, range):
     rng = jtu.rand_default(self.rng())
     _weights = lambda w: abs(w) if weights else None
-    np_fun = lambda a, w: np.histogramdd(a, bins=bins, weights=_weights(w), density=density)
-    jnp_fun = lambda a, w: jnp.histogramdd(a, bins=bins, weights=_weights(w), density=density)
+    np_fun = jtu.ignore_warning(category=RuntimeWarning, message="invalid value.*")(
+        lambda a, w: np.histogramdd(a, bins=bins, weights=_weights(w), density=density, range=range))
+    jnp_fun = lambda a, w: jnp.histogramdd(a, bins=bins, weights=_weights(w), density=density, range=range)
     args_maker = lambda: [rng(shape, dtype), rng((shape[0],), dtype)]
     tol = {jnp.bfloat16: 2E-2, np.float16: 1E-1}
     # np.searchsorted errors on bfloat16 with