anon-iclr-submission/benchname-bug-localization huggingface dataset:
Repository: google/jax
Issue: `jnp.sinc` causes error with NaNs debugging enabled

The [current implementation of `sinc`](https://github.com/google/jax/blob/d39261497c12488632cf7c5d5f19d2c074a565dd/jax/_src/numpy/lax_numpy.py#L1061) function generates NaNs that are detected when `jax_debug_nans` is activated.

**To reproduce:**
```python
import jax
import jax.numpy as jnp
from jax.config import config 
config.update("jax_debug_nans", True)

x = jnp.array([0.])
x = jnp.sinc(x)
```

Output:
```python
---------------------------------------------------------------------------

FloatingPointError                        Traceback (most recent call last)

<ipython-input-5-ca7132c5c117> in <module>()
      6 x = jnp.array([0.])
      7 
----> 8 x = jnp.sinc(x)

[...]

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
    374   if dtypes.issubdtype(xla_shape.element_type(), np.inexact):
    375     if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 376       raise FloatingPointError(f"invalid value (nan) encountered in {name}")
    377     if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):
    378       raise FloatingPointError(f"invalid value (inf) encountered in {name}")

FloatingPointError: invalid value (nan) encountered in div
```


The error comes from the second branch of the `jnp.where` function [here](https://github.com/google/jax/blob/d39261497c12488632cf7c5d5f19d2c074a565dd/jax/_src/numpy/lax_numpy.py#L1068), which has NaNs where `safe_pi_x` is `0`. Those NaNs are filtered out by the `where` function but are nevertheless produced and detected by the NaNs debugger.


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

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 2042408ac..a7622261e 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1064,7 +1064,7 @@ def sinc(x):
   x, = _promote_dtypes_inexact(x)
   eq_zero = lax.eq(x, lax._const(x, 0))
   pi_x = lax.mul(lax._const(x, pi), x)
-  safe_pi_x = where(eq_zero, lax._const(x, 0), pi_x)
+  safe_pi_x = where(eq_zero, lax._const(x, 1), pi_x)
   return where(eq_zero, _sinc_maclaurin(0, pi_x),
                lax.div(lax.sin(safe_pi_x), safe_pi_x))
 
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
index 005de7d59..4aef41cf6 100644
--- a/tests/debug_nans_test.py
+++ b/tests/debug_nans_test.py
@@ -37,6 +37,10 @@ class DebugNaNsTest(jtu.JaxTestCase):
   def tearDown(self):
     config.update("jax_debug_nans", self.cfg)
 
+  def testSinc(self):
+    # Regression test for #6936
+    self.assertEqual(jnp.sinc(0.0), 1.0)
+
   def testSingleResultPrimitiveNoNaN(self):
     A = jnp.array([[1., 2.], [2., 3.]])
     ans = jnp.tanh(A)
Repository: google/jax
Issue: eval_shape appears to hang when tracing a large random sample

Calling

```python
import jax
from jax import random
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11)]))
```

returns `ShapeDtypeStruct(shape=(100000000000,), dtype=float32)`. I'd expect that calling

```python
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11), int(1e11)]))
```

would similarly return `ShapeDtypeStruct(shape=(100000000000,100000000000), dtype=float32)` (of course, if we *actually* ran the function we'd get an OOM error). 

But instead, this call seems to end up in an ~infinite loop (at least, it hangs for > 2 minutes). Reproduction here:
https://colab.research.google.com/drive/14IZF3WUTaq976jaUC5tehKTPwt3YhNTe?usp=sharing

Is this expected?


Changed files:
jax/_src/random.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 4fb534ae5..965c4c120 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -304,14 +304,15 @@ def _random_bits(key, bit_width, shape):
   max_count = int(np.ceil(bit_width * size / 32))
 
   nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
+
   if not nblocks:
     bits = threefry_2x32(key, lax.iota(np.uint32, rem))
   else:
-    *subkeys, last_key = split(key, nblocks + 1)
-    blocks = [threefry_2x32(k, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
-              for k in subkeys]
+    keys = split(key, nblocks + 1)
+    subkeys, last_key = keys[:-1], keys[-1]
+    blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
     last = threefry_2x32(last_key, lax.iota(np.uint32, rem))
-    bits = lax.concatenate(blocks + [last], 0)
+    bits = lax.concatenate([blocks.ravel(), last], 0)
 
   dtype = _UINT_DTYPES[bit_width]
   if bit_width == 64:
Repository: google/jax
Issue: jax.numpy.transpose does not sufficiently validate inputs

For example:
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> x = np.zeros((2, 3, 4))

# jnp.transpose is silent on invalid permutations
>>> jnp.transpose(x, (0, 1, 2, 3))
array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

# numpy transpose raises an error.
>>> np.transpose(x, (0, 1, 2, 3))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in transpose
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 658, in transpose
    return _wrapfunc(a, 'transpose', axes)
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 58, in _wrapfunc
    return bound(*args, **kwds)
ValueError: axes don't match array
```

Changed files:
jax/_src/lax/lax.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index ddc3ef117..64e4300e5 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1216,7 +1216,7 @@ def transpose(operand: Array, permutation: Sequence[int]) -> Array:
   operator.
   """
   permutation = tuple(permutation)
-  if permutation == tuple(range(len(permutation))):
+  if permutation == tuple(range(np.ndim(operand))):
     return operand
   else:
     return transpose_p.bind(operand, permutation=permutation)
Repository: google/jax
Issue: ravel, flatten don't return copies of 1D arrays

The docstring for ravel states:
```
Return a contiguous flattened array.

LAX-backend implementation of :func:`ravel`.

The JAX version of this function will return a copy rather than a view of the input.
```

But this isn't true if the array is already flattened:
```python
import jax.numpy as jnp
x = jnp.array([1, 2, 3])
x.ravel().unsafe_buffer_pointer() == x.unsafe_buffer_pointer() #False
```
Why does this matter? It probably doesn't, but the behavior doesn't match the documentation.

Changed files:
jax/_src/numpy/lax_numpy.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index aa8f902e8..c41e63e79 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1103,7 +1103,8 @@ def _sinc_maclaurin_jvp(k, primals, tangents):
   return _sinc_maclaurin(k, x), _sinc_maclaurin(k + 1, x) * t
 
 _ARRAY_VIEW_DOC = """
-The JAX version of this function will return a copy rather than a view of the input.
+The JAX version of this function may in some cases return a copy rather than a
+view of the input.
 """
 
 @_wraps(np.transpose, lax_description=_ARRAY_VIEW_DOC)
Repository: google/jax
Issue: num replicas not correctly detected with `lax.switch` / `lax.cond`

Thanks to @LenaMartens for the repro!

```python
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=32'

import jax

def identity(x):
  return x

def cond_of_pmap(x):
  y = jax.lax.cond(True, lambda x: jax.pmap(identity)(x), lambda x: jax.pmap(identity)(x), x)
  return y

cond_of_pmap(jax.numpy.zeros((8, 2)))
```

```
AssertionError                            Traceback (most recent call last)
google3/third_party/py/jax/interpreters/xla.py in axis_groups(axis_env, name)
    533   mesh_axes = tuple(unsafe_map(partial(axis_read, axis_env), name))
    534   trailing_size, ragged = divmod(axis_env.nreps, prod(axis_env.sizes))
--> 535   assert not ragged
    536   mesh_spec = axis_env.sizes + (trailing_size,)
    537   return _axis_groups(mesh_spec, mesh_axes)

AssertionError: 
```

The issue was in `core.traverse_jaxpr_params`, which assumed that all subjaxprs appear as values in the params dict, while `lax.switch` (aka `lax.cond`) actually has a param value which is a _tuple_ of jaxprs. Those tuple elements were being ignored!

I'm going to make a quick fix which just flattens params to one level, but I think a better approach here would be to have a primitive-keyed table of rules for traversing subjaxprs.

Changed files:
tests/pmap_test.py, jax/core.py

diff --git a/jax/core.py b/jax/core.py
index 115b0e691..5c25057e4 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -307,11 +307,13 @@ def extract_call_jaxpr(
     return (params["call_jaxpr"], new_params)
 
 
+# TODO(mattjj): replace this approach with a primitive-keyed table of rules
 def traverse_jaxpr_params(f, params):
   """Applies f to each jaxpr parameter and returns a tuple of returned values."""
-  return {name: f(param)
+  return {name: f(p)
           for name, param in params.items()
-          if type(param) in (Jaxpr, ClosedJaxpr)}
+          for p in (param if isinstance(param, (tuple, list)) else [param])
+          if type(p) in (Jaxpr, ClosedJaxpr)}
 
 
 def eval_jaxpr(jaxpr: Jaxpr, consts, *args):
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index 385b539c2..0bc6827ea 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1701,6 +1701,17 @@ class PmapTest(jtu.JaxTestCase):
     out_dtype = func(unused_arg).dtype
     self.assertEqual(out_dtype, dtype)
 
+  def test_num_replicas_with_switch(self):
+    # https://github.com/google/jax/issues/7411
+    def identity(x):
+      return x
+
+    def cond_of_pmap(x):
+      y = lax.cond(True, jax.pmap(identity), jax.pmap(identity), x)
+      return y
+
+    cond_of_pmap(jnp.zeros((xla_bridge.device_count(), 2)))
+
 
 class VmapOfPmapTest(jtu.JaxTestCase):
Repository: google/jax
Issue: Use bazel --version to check the bazel version

Please use `bazel --version` to check the version of Bazel. The currently used variant ignores the bazel-startup-options and actually starts a full bazel server which is a) overkill and b) wrong in case of options like `--output-user-root` and no or readonly home folders

Potential change:
```
try:
  version_output = shell([bazel_path, "--version"])
except subprocess.CalledProcessError:
  return False
match = re.search("(Build label:|bazel) *([0-9\\\\.]+)", version_output)
```

I kept the regexp generic enough so that you can still fall back to the current `bazel version` call and parse both kinds of output

Changed files:
build/build.py

diff --git a/build/build.py b/build/build.py
index dc2750420..c74346563 100755
--- a/build/build.py
+++ b/build/build.py
@@ -209,7 +209,7 @@ def get_bazel_path(bazel_path_flag):
 
 def check_bazel_version(bazel_path):
   try:
-    version_output = shell([bazel_path, "--bazelrc=/dev/null", "version"])
+    version_output = shell([bazel_path, "--version"])
   except subprocess.CalledProcessError:
     return False
   match = re.search("Build label: *([0-9\\\\.]+)[^0-9\\\\.]", version_output)
Repository: google/jax
Issue: jax.scipy.stats.beta.pdf handles x=0, x=1 incorrectly

```
a = b = 1.
x = np.array([0., 1.])
scipy.stats.beta.pdf(x, a, b)
# array([1., 1.])
jax.scipy.stats.beta.pdf(x, a, b)
# DeviceArray([nan, nan], dtype=float32) 
```

Same issue and resolution as #7256 basically--use xlogy/xlog1py here:
https://github.com/google/jax/blob/f6f1debf700b55ad8c0ecc1daa0ed1f677c2bf00/jax/_src/scipy/stats/beta.py#L30-L31



Changed files:
tests/scipy_stats_test.py, jax/_src/scipy/stats/beta.py

diff --git a/jax/_src/scipy/stats/beta.py b/jax/_src/scipy/stats/beta.py
index 64a9e59cb..eb3135de3 100644
--- a/jax/_src/scipy/stats/beta.py
+++ b/jax/_src/scipy/stats/beta.py
@@ -18,7 +18,7 @@ from jax import lax
 from jax._src.numpy.util import _wraps
 from jax._src.numpy.lax_numpy import (_promote_args_inexact, _constant_like,
                                       where, inf, logical_or)
-from jax.scipy.special import betaln
+from jax.scipy.special import betaln, xlogy, xlog1py
 
 
 @_wraps(osp_stats.beta.logpdf, update_doc=False)
@@ -27,8 +27,8 @@ def logpdf(x, a, b, loc=0, scale=1):
   one = _constant_like(x, 1)
   shape_term = lax.neg(betaln(a, b))
   y = lax.div(lax.sub(x, loc), scale)
-  log_linear_term = lax.add(lax.mul(lax.sub(a, one), lax.log(y)),
-                            lax.mul(lax.sub(b, one), lax.log1p(lax.neg(y))))
+  log_linear_term = lax.add(xlogy(lax.sub(a, one), y),
+                            xlog1py(lax.sub(b, one), lax.neg(y)))
   log_probs = lax.sub(lax.add(shape_term, log_linear_term), lax.log(scale))
   return where(logical_or(lax.gt(x, lax.add(loc, scale)),
                           lax.lt(x, loc)), -inf, log_probs)
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
index 548a5e373..154f29b03 100644
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -148,6 +148,13 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
     self._CompileAndCheck(lax_fun, args_maker,
                           rtol={np.float32: 2e-3, np.float64: 1e-4})
 
+  def testBetaLogPdfZero(self):
+    # Regression test for https://github.com/google/jax/issues/7645
+    a = b = 1.
+    x = np.array([0., 1.])
+    self.assertAllClose(
+      osp_stats.beta.pdf(x, a, b), lsp_stats.beta.pdf(x, a, b), atol=1E-6)
+
   @genNamedParametersNArgs(3)
   def testCauchyLogPdf(self, shapes, dtypes):
     rng = jtu.rand_default(self.rng())
Repository: google/jax
Issue: Interaction between `vmap` and new PRNG

### Description

Hi, I'm having trouble seeing how the following MWE using [Diffrax](https://github.com/patrick-kidger/diffrax) errors:
```py
import jax
from jax import numpy as jnp
from diffrax import diffeqsolve, ControlTerm, Euler, MultiTerm, ODETerm, SaveAt, VirtualBrownianTree

B = 20
D = 5

tol = 1e-3
dt0 = 0.05
T0 = 0.0
Tf = 1.0

def f(x0, key):
    drift = lambda t, x, args: jnp.zeros((D, ))
    diffusion = lambda t, x, args: jnp.zeros((D, D))
    brownian_motion = VirtualBrownianTree(T0, Tf, tol=tol, shape=x0.shape, key=key)
    terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))
    solver = Euler()
    saveat = SaveAt(dense=True)
    sol = diffeqsolve(terms, solver, T0, Tf, dt0=dt0, y0=x0, saveat=saveat)
    return sol.evaluate(Tf)

x0 = jnp.zeros((D, ))
key = jax.random.key(1)
print(f(x0, key))

x0 = jnp.zeros((B, D))
key = jax.random.split(key, B)
print(jax.vmap(f)(x0, key))
```
In particular, the last three lines (with `vmap`) cause an error, while the previous three lines seem to work fine.
I raised [this issue in Diffrax](https://github.com/patrick-kidger/diffrax/issues/267) and was told that it was an issue only with the new PRNG, so the bug is in Jax rather than Diffrax. Hence I'm raising it here.

Here are links for:
- [The stack trace.](https://pastebin.com/46hbkCKb)
- [The output of my `conda list`.](https://pastebin.com/3xjrfYYz)

### What jax/jaxlib version are you using?

0.4.12/0.4.12

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10.11, WSL 2 -> Ubuntu 22.04.2 LTS

### NVIDIA GPU info

_No response_

Changed files:
jax/_src/lax/lax.py, tests/random_test.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 6151b48af..c397c9fdb 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3488,7 +3488,6 @@ def _select_hlo_lowering_opaque(ctx, which, *cases):
   avals_in = ctx.avals_in
   aval_out, = ctx.avals_out
   assert all(aval_case == aval_out for aval_case in avals_in[1:])
-  assert avals_in[0].ndim == aval_out.ndim
   select_lower = _select_hlo_lowering
 
   physical_aval_out = core.physical_aval(aval_out)
diff --git a/tests/random_test.py b/tests/random_test.py
index a08ba1b15..c99942519 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1860,6 +1860,15 @@ class KeyArrayTest(jtu.JaxTestCase):
     self.assertIsInstance(ys, random.KeyArray)
     self.assertEqual(ys.shape, (3, 2))
 
+  @skipIf(not config.jax_enable_custom_prng,
+          'requires config.jax_enable_custom_prng')
+  def test_select_scalar_cond(self):
+    # regression test for https://github.com/google/jax/issues/16422
+    ks = self.make_keys(3)
+    ys = lax.select(True, ks, ks)
+    self.assertIsInstance(ys, random.KeyArray)
+    self.assertEqual(ys.shape, (3,))
+
   @skipIf(not config.jax_enable_custom_prng,
           'requires config.jax_enable_custom_prng')
   def test_select2(self):
Repository: google/jax
Issue: jax.random.fold_in function is using static_argnum=(1,)

Currently, the implementation of `jax.random.fold_in` function treats second argument `data` as static:
https://github.com/google/jax/blob/4c34541c00c02fa750b63a6ea9149909e6c4078f/jax/random.py#L193
Is that the intended behavior? 
If I understand correctly, it means, for instance, that innocent-looking loops of sort:
```python
base_key = jax.random.PRNGKey(0)
for i in range(N): 
   do_something_random(jax.random.fold_in(base_key, i))
```
cause fold_in to recompile at every iteration of the loop and take a lot of time. 
On the other hand, if we jit-compile `fold_in` function, things seem to work correctly and fast; like in this gist (thanks @christopherhesse): 
https://gist.github.com/christopherhesse/f493e516b7786533d76c3ef689cb6a45
Should that be the default behavior? 

Changed files:
jax/random.py

diff --git a/jax/random.py b/jax/random.py
index 551cd8b3d..3203cd949 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -190,7 +190,7 @@ def fold_in(key, data):
   """
   return _fold_in(key, data)
 
-@partial(jit, static_argnums=(1,))
+@jit
 def _fold_in(key, data):
   key2 = lax.tie_in(key, PRNGKey(data))
   return threefry_2x32(key, key2)
Repository: google/jax
Issue: `slogdet` sign is incorrect

In the following, numpy disagrees with JAX
```python
mat = np.array([[[-0.01]],
             [[-0.01]]])
print(np.linalg.slogdet(mat))
print(jnp.linalg.slogdet(jnp.array(mat)))
```
=>
```
(array([-1., -1.]), array([-4.60517019, -4.60517019]))
(DeviceArray([1., 1.]), DeviceArray([-4.60517019, -4.60517019]))
```

Changed files:
jax/numpy/linalg.py, tests/linalg_test.py

diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
index aad84415d..be7a96820 100644
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -60,6 +60,7 @@ def svd(a, full_matrices=True, compute_uv=True):
 
 
 @_wraps(onp.linalg.slogdet)
+@jit
 def slogdet(a):
   a = _promote_arg_dtypes(np.asarray(a))
   dtype = lax.dtype(a)
@@ -72,10 +73,10 @@ def slogdet(a):
   is_zero = np.any(diag == np.array(0, dtype=dtype), axis=-1)
   parity = np.count_nonzero(pivot != np.arange(a_shape[-1]), axis=-1)
   if np.iscomplexobj(a):
-    sign = np.prod(diag / np.abs(diag))
+    sign = np.prod(diag / np.abs(diag), axis=-1)
   else:
     sign = np.array(1, dtype=dtype)
-    parity = parity + np.count_nonzero(diag < 0)
+    parity = parity + np.count_nonzero(diag < 0, axis=-1)
   sign = np.where(is_zero,
                   np.array(0, dtype=dtype),
                   sign * np.array(-2 * (parity % 2) + 1, dtype=dtype))
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 95bf853d3..9520908a7 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -99,19 +99,27 @@ class NumpyLinalgTest(jtu.JaxTestCase):
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
-       "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
-       "n": n, "dtype": dtype, "rng": rng}
-      for n in [0, 4, 10, 200]
+       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
+       "shape": shape, "dtype": dtype, "rng": rng}
+      for shape in [(0, 0), (1, 1), (3, 3), (4, 4), (10, 10), (200, 200),
+                    (2, 2, 2), (2, 3, 3), (3, 2, 2)]
       for dtype in float_types + complex_types
       for rng in [jtu.rand_default()]))
-  def testSlogdet(self, n, dtype, rng):
+  def testSlogdet(self, shape, dtype, rng):
     _skip_if_unsupported_type(dtype)
-    args_maker = lambda: [rng((n, n), dtype)]
+    args_maker = lambda: [rng(shape, dtype)]
 
     self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet, args_maker,
                             check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)
 
+  def testIssue1213(self):
+    for n in range(5):
+      mat = np.array([onp.diag(onp.ones([5], dtype=onp.float32))*(-.01)] * 2)
+      args_maker = lambda: [mat]
+      self._CheckAgainstNumpy(onp.linalg.slogdet, np.linalg.slogdet, args_maker,
+                              check_dtypes=True, tol=1e-3)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}".format(
            jtu.format_shape_dtype_string(shape, dtype)),
Repository: google/jax
Issue: jax.numpy.squeeze: error for nonexistent / nonsingleton axis

Numpy's squeeze() function fails with an error if the passed axis is either missing or has a size other than 1. jax.numpy.squeeze() should presumably have the same behaviour.

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

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 737d685bf..9856b3fb3 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -924,15 +924,19 @@ def ravel(a, order="C"):
 
 @_wraps(onp.squeeze)
 def squeeze(a, axis=None):
-  if 1 not in shape(a):
-    return a
+  shape_a = shape(a)
   if axis is None:
-    newshape = [d for d in shape(a) if d != 1]
+    if 1 not in shape_a:
+      return a
+    newshape = [d for d in shape_a if d != 1]
   else:
     if isinstance(axis, int):
       axis = (axis,)
     axis = frozenset(_canonicalize_axis(i, ndim(a)) for i in axis)
-    newshape = [d for i, d in enumerate(shape(a))
+    if _any(shape_a[a] != 1 for a in axis):
+      raise ValueError("cannot select an axis to squeeze out which has size "
+                       "not equal to one")
+    newshape = [d for i, d in enumerate(shape_a)
                 if d != 1 or i not in axis]
   return lax.reshape(a, newshape)
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index fd5b67404..cb5b276d4 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1478,6 +1478,23 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_inshape={}_axis={}".format(
+          jtu.format_shape_dtype_string(arg_shape, dtype), ax),
+       "arg_shape": arg_shape, "dtype": dtype, "ax": ax,
+       "rng_factory": jtu.rand_default}
+      for arg_shape, ax in [
+          ((3,), 0),
+          ((1, 3), 1),
+          ((1, 3, 1), (0, 1))]
+      for dtype in default_dtypes))
+  def testSqueezeFailsOnNonsingletonAxis(self, arg_shape, dtype, ax,
+                                         rng_factory):
+    rng = rng_factory()
+    x = jnp.zeros(arg_shape, dtype=dtype)
+    fun = lambda: jnp.squeeze(x, ax)
+    self.assertRaisesRegex(ValueError, "cannot select an axis to squeeze", fun)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}_weights={}_returned={}".format(
           jtu.format_shape_dtype_string(shape, dtype),
Repository: google/jax
Issue: jnp.einsum docstring missing equation arg

### Description

The documented arg list for einsum is missing the first argument.

https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.einsum.html
`jax.numpy.einsum(*operands, out=None, optimize='optimal', precision=None, _use_xeinsum=False)`
compare to
https://numpy.org/doc/stable/reference/generated/numpy.einsum.html
`numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)`
(noting the 'subscripts' for the first arg)

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

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 b94f8336c..684a280c0 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3112,6 +3112,7 @@ the implementation.
 
 @util._wraps(np.einsum, lax_description=_EINSUM_DOC, skip_params=['out'])
 def einsum(
+    subscripts,
     *operands,
     out=None,
     optimize="optimal",
@@ -3119,6 +3120,7 @@ def einsum(
     _use_xeinsum=False,
     _dot_general=lax.dot_general,
 ):
+  operands = (subscripts, *operands)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.einsum is not supported.")
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 2f7c5ca9c..f4a0a83ae 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -5251,7 +5251,7 @@ class NumpySignaturesTest(jtu.JaxTestCase):
     }
 
     extra_params = {
-      'einsum': ['precision'],
+      'einsum': ['subscripts', 'precision'],
       'einsum_path': ['subscripts'],
       'take_along_axis': ['mode'],
     }
Repository: google/jax
Issue: ModuleNotFoundError: No module named 'requests' when executing `jax.distributed.initialize()` on TPU Pods

### Description

```python
import jax
jax.distributed.initialize()
```

Error:

```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/distributed.py", line 168, in initialize
    global_state.initialize(coordinator_address, num_processes, process_id, local_device_ids)
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/distributed.py", line 47, in initialize
    clusters.ClusterEnv.auto_detect_unset_distributed_params(
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cluster.py", line 55, in auto_detect_unset_distributed_params
    coordinator_address = env.get_coordinator_address()
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 54, in get_coordinator_address
    return cls._get_worker_endpoints()[0].split(':')[2] + ':8476'
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 74, in _get_worker_endpoints
    return get_metadata('worker-network-endpoints').split(',')
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayaka/nfs_share/venv/lib/python3.11/site-packages/jax/_src/clusters/cloud_tpu_cluster.py", line 23, in get_metadata
    import requests  # pytype: disable=import-error
    ^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'requests'
```

I think `requests` should be automatically installed when installing JAX.

### What jax/jaxlib version are you using?

jax 0.4.14

### Which accelerator(s) are you using?

TPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_

Changed files:
setup.py

diff --git a/setup.py b/setup.py
index de67ff9ef..36294aad6 100644
--- a/setup.py
+++ b/setup.py
@@ -85,8 +85,11 @@ setup(
 
         # Cloud TPU VM jaxlib can be installed via:
         # $ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-        'tpu': [f'jaxlib=={_current_jaxlib_version}',
-                f'libtpu-nightly=={_libtpu_version}'],
+        'tpu': [
+          f'jaxlib=={_current_jaxlib_version}',
+          f'libtpu-nightly=={_libtpu_version}',
+          'requests',  # necessary for jax.distributed.initialize
+        ],
 
         # $ pip install jax[australis]
         'australis': ['protobuf>=3.13,<4'],