internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/jax

Issue:
`jnp.squeeze` will fail unexpectedly with `jit` compilation
`jnp.squeeze` will fail unexpectedly with `jit` compilation but direct invocation will succeed. For example,

```py
import jax

def fn(inp):
    axis = jax.numpy.array(1, dtype=jax.numpy.int32)
    return jax.numpy.squeeze(inp, axis)

mykey = jax.random.PRNGKey(0)
inp = jax.random.uniform(mykey, [3, 1], jax.numpy.float32, minval=0, maxval=1)

print(fn(inp))
# [0.9653214  0.31468165 0.63302994]

jax.jit(fn)(inp)
# TypeError: iteration over a 0-d array
```

I think the input is valid since the dim-`1` has shape 1.

## Version

jax 0.3.17


Discussion:
BTW, `jnp.expand_dims` has a similar issue

```
import jax

def fn(inp):
    axis = jax.numpy.array(0, dtype=jax.numpy.int64)
    return jax.numpy.expand_dims(inp, axis)

inp = jax.numpy.array(128, dtype=jax.numpy.float32)

print(fn(inp))
# [128.]

jax.jit(fn)(inp)
# TypeError: iteration over a 0-d array
```
The error is misleading here, but I believe the issue is that in both cases the `axis` argument must be static, and when you pass a jax array to the function within `jit` it is no longer static. We should update the `_ensure_index_tuple` helper function to return a more suitable error in this case.

Resolving patch:
diff --git a/jax/_src/api_util.py b/jax/_src/api_util.py
--- a/jax/_src/api_util.py
+++ b/jax/_src/api_util.py
@@ -37,6 +37,7 @@
 
 def _ensure_index(x: Any) -> Union[int, Tuple[int, ...]]:
   """Ensure x is either an index or a tuple of indices."""
+  x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
   try:
     return operator.index(x)
   except TypeError:
@@ -44,6 +45,7 @@ def _ensure_index(x: Any) -> Union[int, Tuple[int, ...]]:
 
 def _ensure_index_tuple(x: Any) -> Tuple[int, ...]:
   """Convert x to a tuple of indices."""
+  x = core.concrete_or_error(None, x, "expected a static index or sequence of indices.")
   try:
     return (operator.index(x),)
   except TypeError:


Test patch:
Repository: google/jax

Issue:
[x64] Reductions are too eager to promote to 64-bit
```python
In [1]: import jax.numpy as jnp

In [2]: jnp.arange(4, dtype='int8').sum()
Out[2]: DeviceArray(6, dtype=int64)
```

The root cause is this line: https://github.com/google/jax/blob/d849f495193c69e8c133d0c6a424ba80d72938cb/jax/_src/numpy/reductions.py#L81

It's not an issue for floating-point types, because numpy does not promote those in most cases


Discussion:
Still need to flip the default for this to be fixed

Resolving patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -86,17 +86,21 @@ def _reduction(a, name, np_fun, op, init_val, has_identity=True,
     if not _all(core.greater_equal_dim(shape[d], 1) for d in pos_dims):
       raise ValueError(f"zero-size array to reduction operation {name} which has no identity")
 
-  result_dtype = dtypes.canonicalize_dtype(dtype or dtypes.dtype(a))
+  result_dtype = dtype or dtypes.dtype(a)
 
-  # promote_integers=True matches NumPy's behavior for sum() and prod(), which promotes
-  # all int-like inputs to the widest available dtype.
   if dtype is None and promote_integers:
+    # Note: NumPy always promotes to 64-bit; jax instead promotes to the
+    # default dtype as defined by dtypes.int_ or dtypes.uint.
     if dtypes.issubdtype(result_dtype, np.bool_):
-      result_dtype = dtypes.canonicalize_dtype(np.int64)
+      result_dtype = dtypes.int_
     elif dtypes.issubdtype(result_dtype, np.unsignedinteger):
-      result_dtype = dtypes.canonicalize_dtype(np.uint64)
+      if np.iinfo(result_dtype).bits < np.iinfo(dtypes.uint).bits:
+        result_dtype = dtypes.uint
     elif dtypes.issubdtype(result_dtype, np.integer):
-      result_dtype = dtypes.canonicalize_dtype(np.int64)
+      if np.iinfo(result_dtype).bits < np.iinfo(dtypes.int_).bits:
+        result_dtype = dtypes.int_
+
+  result_dtype = dtypes.canonicalize_dtype(result_dtype)
 
   if upcast_f16_for_computation and dtypes.issubdtype(result_dtype, np.inexact):
     computation_dtype = _upcast_f16(result_dtype)


Test patch:
Repository: google/jax

Issue:
Array typing/IDE errors since version 0.3.18
### Description

JAX 0.3.18 seems to have included some updates to the way array types are handled (via the introduction of `jax.Array`).

This seems to have broken some things for those of us running type checkers. Basic operations like `some_array.shape` now result in errors in both `pyright` and `mypy`:

![image](https://user-images.githubusercontent.com/6992947/192605080-591e182a-ed32-4ee0-90d8-91b8aae0d155.png)

Tab completion for common array operations like `.shape`, `.reshape()`, `.astype()`, etc also seems broken.

Is this a bug? If so is there a fix planned?

Thanks so much!!

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

jax v0.3.18, v0.3.19

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

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report, and sorry for this issue! In 0.3.18 we introduced more restrictive array types as part of the work in #12049. Those array attributes are defined in the stub file here: https://github.com/google/jax/blob/main/jax/_src/basearray.pyi

This is enough to satisfy mypy in our github CI checks, as well as google's internal fork of pytype, which is used for Google's internal static type checking.

I wonder why it's not working for you? Perhaps the `.pyi` file is somehow not being properly registered?

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,7 @@ def generate_proto(source):
     author='JAX team',
     author_email='jax-dev@google.com',
     packages=find_packages(exclude=["examples"]),
-    package_data={'jax': ['py.typed']},
+    package_data={'jax': ['py.typed', "*.pyi", "**/*.pyi"]},
     python_requires='>=3.7',
     install_requires=[
         'absl-py',


Test patch:
Repository: google/jax

Issue:
PYTHONOPTIMIZE=1 causes KeyErrors
### Description

When I set the flag: `export PYTHONOPTIMIZE=1`, the code:
```
import jax.numpy as jnp
m = jnp.eye(5)
a = jnp.ones(5)
jnp.linalg.solve(m,a)
```
returns the following error:
```
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "error.py", line 6, in <module>
    jnp.linalg.solve(m,a)
  File "/home/ben/.../python3.8/site-packages/jax/_src/numpy/linalg.py", line 553, in solve
    return lax_linalg._solve(a, b)
KeyError: 139645559661888
```

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

jax v0.3.21, jaxlib v0.3.20

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

CPU

### Additional system info

Ubuntu, jax is installed into a site_packages repo for ROS.

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! I see the issue; there's an assert statement we're counting on running.

Resolving patch:
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -898,7 +898,8 @@ def get_atom(t: JaxprTracer) -> Atom:
 
   def newvar(t: JaxprTracer) -> Var:
     var = gensym(type_substitute(t.aval))
-    assert t_to_var.setdefault(id(t), var) is var
+    var_ = t_to_var.setdefault(id(t), var)
+    assert var is var_
     return var
 
   def type_substitute(aval: AbstractValue) -> AbstractValue:


Test patch:
Repository: google/jax

Issue:
jax.scipy.special.erf: Doesn't support complex arguments despite documentation promising that
### Description

From the documentation ( https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.special.erf.html ):
```
Returns the error function of complex argument.
```

But:
```
from jax.scipy.special import erf
erf(1j)
```

gives 
```
TypeError: erf does not accept dtype complex64. Accepted dtypes are subtypes of floating.
```

.

So either the implementation or the documentation seem broken here.

Needless to say, it would be very useful if `erf` would indeed support complex arguments. A related feature request: #9098

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

jax v0.3.21, jaxlib v0.3.20

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

CPU

### Additional system info

Python 3.10.6, linux

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! It's a bit confusing, but the documentation comes from scipy rather than JAX (notice the "*Original docstring below*" on the page you linked to).

We should add a note to the docs mentioning the restriction on input types.

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -69,7 +69,8 @@ def gammaincc(a, x):
   return lax.igammac(a, x)
 
 
-@_wraps(osp_special.erf, module='scipy.special')
+@_wraps(osp_special.erf, module='scipy.special', skip_params=["out"],
+        lax_description="Note that the JAX version does not support complex inputs.")
 def erf(x):
   x, = _promote_args_inexact("erf", x)
   return lax.erf(x)


Test patch:
Repository: google/jax

Issue:
jax.numpy.linalg.lstsq throws IndexError with zero rows
### Description

[jax.numpy.linalg.lstsq](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.lstsq.html) throws an IndexError when the number of rows is zero:

```python3
import numpy as np
import jax.numpy as jnp

A = np.zeros([0, 3])
b = np.zeros([0])
x, residuals, rank, s = np.linalg.lstsq(A, b, None)
print(x) # [0. 0. 0.]

A = jnp.zeros([0, 3])
b = jnp.zeros([0])
x, residuals, rank, s = jnp.linalg.lstsq(A, b, None)
# IndexError: index is out of bounds for axis 0 with size 0
```

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

jax 0.3.21, jaxlib 0.3.20

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

CPU

### Additional system info

Python 3.10.7, macOS 11.7

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report!

Resolving patch:
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -601,17 +601,22 @@ def _lstsq(a: ArrayLike, b: ArrayLike, rcond: Optional[float], *,
       f"{b.ndim}-dimensional array given. Array must be one or two-dimensional")
   m, n = a.shape
   dtype = a.dtype
-  if rcond is None:
-    rcond = jnp.finfo(dtype).eps * max(n, m)
+  if a.size == 0:
+    s = jnp.empty(0, dtype=a.dtype)
+    rank = jnp.array(0, dtype=int)
+    x = jnp.empty((n, *b.shape[1:]), dtype=a.dtype)
   else:
-    rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
-  u, s, vt = svd(a, full_matrices=False)
-  mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
-  rank = mask.sum()
-  safe_s = jnp.where(mask, s, 1).astype(a.dtype)
-  s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
-  uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
-  x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
+    if rcond is None:
+      rcond = jnp.finfo(dtype).eps * max(n, m)
+    else:
+      rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
+    u, s, vt = svd(a, full_matrices=False)
+    mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
+    rank = mask.sum()
+    safe_s = jnp.where(mask, s, 1).astype(a.dtype)
+    s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
+    uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
+    x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
   # Numpy returns empty residuals in some cases. To allow compilation, we
   # default to returning full residuals in all cases.
   if numpy_resid and (rank < n or m <= n):


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -890,6 +890,9 @@ def testMultiDot(self, shapes, dtype):
           ((4, 6), (4,)),
           ((6, 6), (6, 1)),
           ((8, 6), (8, 4)),
+          ((0, 3), (0,)),
+          ((3, 0), (3,)),
+          ((3, 1), (3, 0)),
       ]
     ],
     rcond=[-1, None, 0.5],
Repository: google/jax

Issue:
jet of dynamic slice does not match jet of gather
This is blocking #12219

For example:
```python
import jax.numpy as jnp
from jax import lax
from jax.experimental.jet import jet

f_gather = lambda x: x[1:]
f_dynamic_slice = lambda x: lax.dynamic_slice(x, (1,), (len(x) - 1,))

# These two functions by design have identical behavior for 1D inputs:
x = jnp.arange(4.0)
print(f_gather(x))
# [1. 2. 3.]
print(f_dynamic_slice(x))
# [1. 2. 3.]

# However, their jet transforms do not match
h = 0.1 * x
print(jet(f_gather, (x,), ((h,),)))
# (DeviceArray([1., 2., 3.], dtype=float32), [DeviceArray([0.1, 0.2, 0.3], dtype=float32)])
print(jet(f_dynamic_slice, (x,), ((h,),)))
# (DeviceArray([1., 2., 3.], dtype=float32), [DeviceArray([0. , 0.1, 0.2], dtype=float32)])
```
It looks like the jet rule for dynamic_slice is not behaving correctly.


Discussion:
I think it just needs a (fairly) trivial `jet` rule of its own. `dynamic_slice` doesn't like having its `index` zeroed out by the `deflinear` logic.

Resolving patch:
diff --git a/jax/experimental/jet.py b/jax/experimental/jet.py
--- a/jax/experimental/jet.py
+++ b/jax/experimental/jet.py
@@ -329,12 +329,20 @@ def linear_prop(prim, primals_in, series_in, **params):
 deflinear(lax.rev_p)
 deflinear(lax.transpose_p)
 deflinear(lax.slice_p)
-deflinear(lax.dynamic_slice_p)
 deflinear(lax.reduce_sum_p)
 deflinear(lax.reduce_window_sum_p)
 deflinear(lax.fft_p)
 deflinear(dispatch.device_put_p)
 
+def _dynamic_slice_jet_rule(primals_in, series_in, **params):
+  operand, *start_indices = primals_in
+  primal_out = lax.dynamic_slice_p.bind(operand, *start_indices, **params)
+  series_out = [lax.dynamic_slice_p.bind(terms_in[0], *start_indices, **params)
+                for terms_in in zip(*series_in)]
+  return primal_out, series_out
+
+jet_rules[lax.dynamic_slice_p] = _dynamic_slice_jet_rule
+
 def _dynamic_update_slice_jet_rule(primals_in, series_in, **params):
   operand, update, *start_indices = primals_in
   primal_out = lax.dynamic_update_slice_p.bind(operand, update, *start_indices)


Test patch:
diff --git a/tests/jet_test.py b/tests/jet_test.py
--- a/tests/jet_test.py
+++ b/tests/jet_test.py
@@ -311,7 +311,7 @@ def test_cummax(self):     self.unary_check(partial(lax.cummax, axis=0))
   @jtu.skip_on_devices("tpu")
   def test_cummin(self):     self.unary_check(partial(lax.cummin, axis=0))
   @jtu.skip_on_devices("tpu")
-  def test_dynamic_slice(self):     self.unary_check(partial(lax.dynamic_slice, start_indices=(0,0), slice_sizes=(1,1)))
+  def test_dynamic_slice(self): self.unary_check(partial(lax.dynamic_slice, start_indices=(1,2), slice_sizes=(1,1)))
   @jtu.skip_on_devices("tpu")
   def test_dynamic_update_slice(self): self.unary_check(partial(lax.dynamic_update_slice, start_indices=(1,2), update=jnp.arange(6.0).reshape(2, 3)))
Repository: google/jax

Issue:
jax.random.categorical doesn't respect shape parameter with broadcasted logits
### Description

The docs for `jax.random.categorical` say that the output of `jax.random.categorical` is `shape` when `shape` is provided. However, when provided `logits` which have to be broadcasted with `shape`, it just ignores the shape parameter after checking that it's broadcast compatible. I believe this can be fixed by broadcasting logits appropriately.

```
import jax
out = jax.random.categorical(jax.random.PRNGKey(0), logits=jax.numpy.ones( (4, 1, 10)), axis=-1, shape=(4, 8))

# out.shape is (4, 1), should be (4, 8)

assert out.shape == (4, 8), f"out.shape is {out.shape}, should be (4, 8)"
```



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

0.3.23

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

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1327,10 +1327,12 @@ def categorical(key: KeyArray,
     shape = tuple(shape)
     _check_shape("categorical", shape, batch_shape)
 
-  sample_shape = shape[:len(shape)-len(batch_shape)]
+  shape_prefix = shape[:len(shape)-len(batch_shape)]
+  logits_shape = list(shape[len(shape) - len(batch_shape):])
+  logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])
   return jnp.argmax(
-      gumbel(key, sample_shape + logits_arr.shape, logits_arr.dtype) +
-      lax.expand_dims(logits_arr, tuple(range(len(sample_shape)))),
+      gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype) +
+      lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),
       axis=axis)
 
 


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1429,6 +1429,19 @@ def f():
     # just lower, don't run, takes too long
     jax.jit(f).lower()
 
+  @jtu.sample_product(shape=[(3, 4)],
+                      logits_shape_base=[(3, 4), (3, 1), (1, 4)],
+                      axis=[-3, -2, -1, 0, 1, 2])
+  def test_categorical_shape_argument(self, shape, logits_shape_base, axis):
+    # https://github.com/google/jax/issues/13124
+    logits_shape = list(logits_shape_base)
+    logits_shape.insert(axis % (len(logits_shape_base) + 1), 10)
+    assert logits_shape[axis] == 10
+    logits = jnp.ones(logits_shape)
+    samples = jax.random.categorical(jax.random.PRNGKey(0), logits=logits,
+                                     axis=axis, shape=shape)
+    self.assertEqual(samples.shape, shape)
+
 
 class KeyArrayTest(jtu.JaxTestCase):
   # Key arrays involve:
Repository: google/jax

Issue:
[sparse] bcoo_extract incorrect for zero sparse dimensions
Found in the course of #13624 

Minimal repro:
```python
import jax.numpy as jnp
from jax.experimental import sparse

mat = jnp.zeros(3, dtype='float32')
indices = jnp.zeros((3, 2, 0), dtype='int32')

jax.jit(sparse.bcoo_extract)(indices, mat)  # Error
```


Discussion:

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -400,7 +400,7 @@ def bcoo_extract(indices: Array, mat: Array) -> Array:
 @bcoo_extract_p.def_impl
 def _bcoo_extract_impl(indices, mat):
   mat = jnp.asarray(mat)
-  n_batch, n_sparse, _, _ = _validate_bcoo_indices(indices, mat.shape)
+  n_batch, n_sparse, _, nse = _validate_bcoo_indices(indices, mat.shape)
 
   ind_slices = tuple(np.zeros(s, int) if i_s == 1 else np.arange(s)
                      for s, i_s in zip(mat.shape[:n_batch], indices.shape[:n_batch]))
@@ -412,8 +412,13 @@ def _bcoo_extract_impl(indices, mat):
   batch_ind = tuple(grid)[:-1]
 
   if not sparse_ind + batch_ind:
-    return mat[None]
-  return mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+    result = mat[None]
+  else:
+    result = mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+  if n_sparse == 0 and nse != 1:
+    result = lax.broadcast_in_dim(
+      result, _tuple_replace(result.shape, n_batch, nse), range(result.ndim))
+  return result
 
 @bcoo_extract_p.def_abstract_eval
 def _bcoo_extract_abstract_eval(indices, mat):


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -904,6 +904,17 @@ def test_bcoo_extract_ad(self, shape, dtype, n_batch, n_dense):
     self.assertEqual(j1.shape, data.shape + M.shape)
     self.assertEqual(hess.shape, data.shape + 2 * M.shape)
 
+  def test_bcoo_extract_zero_nse(self):
+    # Regression test for https://github.com/google/jax/issues/13653
+
+    # (n_batch, n_sparse, n_dense) = (1, 0, 0), nse = 2
+    args_maker = lambda: (jnp.zeros((3, 2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
+    # (n_batch, n_sparse, n_dense) = (0, 0, 1), nse = 2
+    args_maker = lambda: (jnp.zeros((2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
   @jtu.sample_product(
     [dict(shape=shape, n_batch=n_batch, n_dense=n_dense)
       for shape in [(), (5,), (5, 8), (3, 4, 5), (3, 4, 3, 2)]
Repository: google/jax

Issue:
Zero-length indices in `jax.numpy.mgrid`
### Description

```python
np.mgrid[()]  # yields an integer type array of zero length
jnp.mgrid[()]  # raises a ValueError because it tries to stack a tuple of zero length
```

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

jax commit f729da4a368e9, jaxlib 0.4.1

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

CPU/GPU

### Additional system info

python 3.10.9, Linux (Arch Linux)

### NVIDIA GPU info

```
| NVIDIA-SMI 525.60.13    Driver Version: 525.60.13    CUDA Version: 12.0     |              
|-------------------------------+----------------------+----------------------+              
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |              
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |              
|                               |                      |               MIG M. |              
|===============================+======================+======================|              
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |              
|  0%   40C    P8    25W / 350W |     17MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
```


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/index_tricks.py b/jax/_src/numpy/index_tricks.py
--- a/jax/_src/numpy/index_tricks.py
+++ b/jax/_src/numpy/index_tricks.py
@@ -54,7 +54,11 @@ def __getitem__(self, key: Union[slice, Tuple[slice, ...]]) -> Union[Array, List
     with jax.numpy_dtype_promotion('standard'):
       output = _promote_dtypes(*output)
     output_arr = jnp.meshgrid(*output, indexing='ij', sparse=self.sparse)
-    return output_arr if self.sparse else jnp.stack(output_arr, 0)
+    if self.sparse:
+      return output_arr
+    if len(output_arr) == 0:
+      return jnp.arange(0)
+    return jnp.stack(output_arr, 0)
 
 
 class _Mgrid(_IndexGrid):


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
@@ -4353,6 +4353,7 @@ def testMgrid(self):
     # wrap indexer for appropriate dtype defaults.
     np_mgrid = _indexer_with_default_outputs(np.mgrid)
     assertAllEqual = partial(self.assertAllClose, atol=0, rtol=0)
+    assertAllEqual(np_mgrid[()], jnp.mgrid[()])
     assertAllEqual(np_mgrid[:4], jnp.mgrid[:4])
     assertAllEqual(np_mgrid[:4,], jnp.mgrid[:4,])
     assertAllEqual(np_mgrid[:4], jax.jit(lambda: jnp.mgrid[:4])())
Repository: google/jax

Issue:
Typo in function `initial_step_size` in `jax/experimental/ode.py`
### Description

The function `initial_step_size` is based from an algorithm in [this book](http://mezbanhabibi.ir/wp-content/uploads/2020/01/ordinary-differential-equations-vol.1.-Nonstiff-problems.pdf) (page 169). This is the algorithm:
<img width="486" alt="image" src="https://user-images.githubusercontent.com/49232747/212095806-ce29dcd6-40f1-4bc1-af54-f567330c0acb.png">

In step `e`, we need to compute `max(d1, d2)`, but in the `initial_step_size` function you compute `max(d1 + d2)`.

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

jax v0.4.1, jaxlib v0.4.1

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

CPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks for noticing this and bringing it to our attention! Amazing!

Resolving patch:
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
--- a/jax/experimental/ode.py
+++ b/jax/experimental/ode.py
@@ -90,7 +90,7 @@ def initial_step_size(fun, t0, y0, order, rtol, atol, f0):
 
   h1 = jnp.where((d1 <= 1e-15) & (d2 <= 1e-15),
                 jnp.maximum(1e-6, h0 * 1e-3),
-                (0.01 / jnp.max(d1 + d2)) ** (1. / (order + 1.)))
+                (0.01 / jnp.maximum(d1, d2)) ** (1. / (order + 1.)))
 
   return jnp.minimum(100. * h0, h1)
 


Test patch: