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

Issue:
Bug in shape checking rule when calling jax.lax.reduce_window
I stumbled on this by accident and the error message says to report it:

```
jax.lax.reduce_window(
  np.zeros((10, 28, 28, 3), dtype=np.float32), jnp.inf, lambda *args: [jax.lax.min(*args)], [1,3,3,1], [1,2,2,1], "SAME")
```

The computation argument appears to be the one triggering the bug, the plain jax.lax.min works fine.


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1147,7 +1147,15 @@ def reduce(operands: Array, init_values: Array, computation: Callable,
 @cache()
 def _reduction_jaxpr(computation, aval):
   pval = pe.PartialVal.unknown(aval)
-  comp = lu.wrap_init(lambda x, y: (computation(x, y),))
+  @lu.wrap_init
+  def comp(x, y):
+    result = computation(x, y)
+    if not (isinstance(result, core.Tracer) or core.valid_jaxtype(result)):
+      raise ValueError(
+          f"Invalid return type from reduction function: {type(result)}\n"
+          f"Reduction functions should only return an array.\n"
+          f"Full return value: {result}")
+    return (result,)
   jaxpr, _, consts = pe.trace_to_jaxpr(comp, (pval, pval), instantiate=False)
   return jaxpr, consts
 
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
@@ -517,6 +517,10 @@ def trace_to_subjaxpr(main: core.MainTrace, instantiate: Union[bool, Sequence[bo
   trace = JaxprTrace(main, core.cur_sublevel())
   in_tracers = map(trace.new_arg, pvals)
   ans = yield in_tracers, {}
+  assert isinstance(ans, (list, tuple)), (
+      f"Got unexpected return type when tracing function to jaxpr: {ans}")
+  assert all(isinstance(x, core.Tracer) or core.valid_jaxtype(x) for x in ans), (
+      f"Got unexpected return type when tracing function to jaxpr: {ans}")
   instantiate = [instantiate] * len(ans) if isinstance(instantiate, bool) else instantiate
   out_tracers = map(trace.full_raise, map(core.full_lower, ans))
   out_tracers = map(partial(instantiate_const_at, trace), instantiate, out_tracers)


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1571,6 +1571,15 @@ def zero_stride_test():
       with self.assertRaisesRegex(TypeError, "must have every element be"):
         failure_fun()
 
+    with self.assertRaisesRegex(
+        ValueError,
+        "Invalid return type from reduction function: <class 'list'>\n"
+        "Reduction functions should only return an array.\n"
+        "Full return value: .*"):
+      return lax.reduce_window(
+          np.ones((1,)), 0., lambda x, y: [x + y],
+          padding='VALID', window_dimensions=(1,), window_strides=(1,))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": (f"_shape={shape}_windowdimensions={window_dimensions}"
                          f"_basedilation={base_dilation}_windowdilation="
Repository: google/jax

Issue:
bincount of empty array raises
Running
`jnp.bincount(jnp.int32([]))`
raises
`ValueError: zero-size array to reduction operation max which has no identity`

In contrast, `np.bincount([])` produces `array([], dtype=int64)`


Discussion:
Thanks for the report!

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
@@ -1659,7 +1659,7 @@ def bincount(x, weights=None, minlength=0, *, length=None):
     x = core.concrete_or_error(asarray, x,
       "The error occured because of argument 'x' of jnp.bincount. "
       "To avoid this error, pass a static `length` argument.")
-    length = max(x) + 1
+    length = max(x, initial=-1) + 1
   length = _max(length, minlength)
   if ndim(x) != 1:
     raise ValueError("only 1-dimensional input supported.")


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
@@ -4001,7 +4001,7 @@ def testAtLeastNdLiterals(self, pytype, dtype, op):
       "weights": weights,
       "minlength": minlength,
       "length": length}
-    for shape in [(5,), (10,)]
+    for shape in [(0,), (5,), (10,)]
     for dtype in int_dtypes
     for weights in [True, False]
     for minlength in [0, 20]
Repository: google/jax

Issue:
jax.devices() behaves unintuitively
Hi! I noticed some unexpected (?) behaviour in the following code:
```python 
import jax
from jax.config import config
config.update("jax_backend_target", 'grpc://some endpoint:8470')
config.update("jax_xla_backend", 'tpu_driver')
print(jax.devices('cpu')) #prints tpu devices instead of cpu
``` 
I can get the expected behaviour if I do 

```python 
import jax
print(jax.devices('cpu')) #prints cpu devices 
from jax.config import config
config.update("jax_backend_target", 'grpc://11.111.11.1118470')
config.update("jax_xla_backend", 'tpu_driver')
print(jax.devices('cpu')) #now prints cpu devices
print(jax.devices('tpu')) #now prints tpu devices
``` 
I think this may be related the caching of `jax.lib.xla_bridge.get_backend()`. Not sure if this is expected behaviour or a bug. 

I noticed this because I was trying `jit` a few smaller functions on the host vm cpu during a larger TPU computation . I tried using the `backend=cpu` argument and `device_put`, but was unable to obtain the desired behaviour. In the end the only thing that seemed to work was to clear the cache of `get_backend()` reconfigure `jax.config` to cpu.


Discussion:
@jacksonwb
Aren't you two already using the new Cloud TPU alpha program? That doesn't have this problem :)

This is indeed a bug in the `tpu_driver` client. Given that we're focusing our efforts on the new Cloud TPU architecture and you have a workaround, we probably won't prioritize fixing this. Please let me know if I'm underestimating the impact of this though and we can reconsider!
We are using the new Cloud TPU alpha program, but still have some applications on the old Cloud TPU setup due to requirements for GKE support. 

Currently we have a workaround involving calls to `jax.lib.xla_bridge.get_backend.cache_clear()` when switching from TPU to CPU jitting, which seems to be working, but I'm not sure what other unintended consequences this may be creating.

Resolving patch:
diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -144,7 +144,9 @@ def _get_local_backend(platform=None):
 _tpu_backend = None
 
 def _get_tpu_driver_backend(platform):
-  del platform
+  if platform == "cpu":
+    return _get_local_backend("cpu")
+
   global _tpu_backend
   if _tpu_backend is None:
     backend_target = FLAGS.jax_backend_target


Test patch:
Repository: google/jax

Issue:
Bug in jnp.take with negative indexing
Hello, I am seeing the following issue with jax.numpy.take and -1 index: 

`A = jnp.array([[1, 2], [3, 4], [5, 6]])`
`A[0]` --> [1, 2]
`A[-1]` --> [5, 6]
`A.take([0, -1], axis=0)` --> [[1, 2], [1, 2]]
`jnp.take(A, [0, -1], axis=0)` --> [[1, 2], [1, 2]]

Thanks!


Discussion:
Thanks for the report! It looks like the issue is this is lowered to `lax.gather`, which doesn't support negative indices.

One workaround currently would be to use `mode='wrap'`, which does wrap negative indices correctly.

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
@@ -4142,12 +4142,15 @@ def take(a, indices, axis: Optional[int] = None, out=None, mode=None):
   else:
     axis_idx = _canonicalize_axis(axis, ndim(a))
 
-  if mode == "raise":
+  if mode is None:
+    # lax.gather() does not support negative indices, so we wrap them here
+    indices = where(indices < 0, indices + a.shape[axis_idx], indices)
+  elif mode == "raise":
     # TODO(phawkins): we have no way to report out of bounds errors yet.
     raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
   elif mode == "wrap":
     indices = mod(indices, _constant_like(indices, a.shape[axis_idx]))
-  elif mode != "clip" and mode is not None:
+  elif mode != "clip":
     raise ValueError("Invalid mode '{}' for np.take".format(mode))
 
   index_dims = len(shape(indices))


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
@@ -3621,7 +3621,7 @@ def testUnpackbits(self, shape, dtype, axis, bitorder, count):
                                   [cast(Optional[int], None)])
       for dtype in all_dtypes
       for index_dtype in int_dtypes
-      for mode in ['wrap', 'clip']))
+      for mode in [None, 'wrap', 'clip']))
   def testTake(self, shape, dtype, index_shape, index_dtype, axis, mode):
     def args_maker():
       x = rng(shape, dtype)
@@ -3629,7 +3629,10 @@ def args_maker():
       return x, i
 
     rng = jtu.rand_default(self.rng())
-    rng_indices = jtu.rand_int(self.rng(), -5, 5)
+    if mode is None:
+      rng_indices = jtu.rand_int(self.rng(), -shape[axis or 0], shape[axis or 0])
+    else:
+      rng_indices = jtu.rand_int(self.rng(), -5, 5)
     jnp_op = lambda x, i: jnp.take(x, i, axis=axis, mode=mode)
     np_op = lambda x, i: np.take(x, i, axis=axis, mode=mode)
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
Repository: google/jax

Issue:
bad error message when with grad(f, has_aux=True) when f doesn't return a pair
```python
import jax
jax.grad(lambda x: (1, 2, 3), has_aux=True)(1.)
```

```
ValueError: too many values to unpack (expected 2)
```


Discussion:

Resolving patch:
diff --git a/jax/api_util.py b/jax/api_util.py
--- a/jax/api_util.py
+++ b/jax/api_util.py
@@ -75,7 +75,11 @@ def apply_flat_fun_nokwargs(fun, io_tree, py_args):
 @lu.transformation_with_aux
 def flatten_fun_nokwargs2(in_tree, *args_flat):
   py_args = tree_unflatten(in_tree, args_flat)
-  ans, aux = yield py_args, {}
+  pair = yield py_args, {}
+  if not isinstance(pair, (list, tuple)) or len(pair) != 2:
+    raise TypeError("expected function with aux output to return a two-element "
+                    f"tuple, but got type {type(pair)} with value {repr(pair)}")
+  ans, aux = pair
   ans_flat, ans_tree = tree_flatten(ans)
   aux_flat, aux_tree = tree_flatten(aux)
   yield (ans_flat, aux_flat), (ans_tree, aux_tree)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -876,6 +876,16 @@ def test_grad_and_aux_basic(self):
     self.assertAllClose(g, grad(lambda x: x**3)(3.))
     self.assertAllClose(aux, [9.], check_dtypes=False)
 
+  def test_grad_and_aux_error(self):
+    with self.assertRaisesRegex(TypeError, "two-element tuple"):
+      grad(lambda x: (1, 2, 3), has_aux=True)(1.)
+
+    with self.assertRaisesRegex(TypeError, "two-element tuple"):
+      grad(lambda x: x, has_aux=True)(1.)
+
+    with self.assertRaisesRegex(TypeError, "two-element tuple"):
+      grad(lambda x: (x,), has_aux=True)(1.)
+
   def test_grad_and_aux_nested(self):
     def f(x):
       g, aux = grad(lambda x: (x**3, [x**3]), has_aux=True)(x)
@@ -2319,6 +2329,7 @@ def __jax_array__(self):
       self.assertEqual(f(x), f(a))
 
 
+
 class RematTest(jtu.JaxTestCase):
 
   def test_remat_basic(self):
Repository: google/jax

Issue:
Bug in logsumexp when b argument has 0s
It looks like logsumexp can give the wrong answer if it's b argument has 0s in the indices where the a argument is large.  For example
``` python
import jax.numpy as jnp
from jax.scipy.special import logsumexp

a = jnp.array([-1000.0, -2.0])
b = jnp.array([1.0, 0.0])
out = logsumexp(a, b=b) # DeviceArray(-inf, dtype=float32)

# out should be DeviceArray(-1000.0, dtype=float32)
```
The bug seems to happen because the implementation subtracts the max value of a without accounting for b.  A simple fix should be to add these 2 lines of code in the first check if b is not None:
``` python
a = a + jnp.where(b, jnp.log(jnp.abs(b)), -jnp.inf)
b = jnp.sign(b)
```


Discussion:
Thanks for the report! Are you interested in contributing a pull request to fix this issue? 
Sure!
Great! Feel free to submit a pull request with the fix, and ping me when you do. We'll want to make sure to add a new test case that covers this bug (i.e. one that would fail with the current code, but passes with the fix). Let me know if any questions come up along the way

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
@@ -101,7 +101,8 @@ def expit(x):
 @_wraps(osp_special.logsumexp)
 def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
   if b is not None:
-    a, b = jnp.broadcast_arrays(a, b)
+    a, b = _promote_args_inexact("logsumexp", a, b)
+    a = jnp.where(b != 0, a, -jnp.inf)
   pos_dims, dims = _reduction_dims(a, axis)
   amax = jnp.max(a, axis=dims, keepdims=keepdims)
   amax = lax.stop_gradient(lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0)))


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -146,6 +146,14 @@ def lax_fun(array_to_reduce):
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
     self._CompileAndCheck(lax_fun, args_maker)
 
+  def testLogSumExpZeros(self):
+    # Regression test for https://github.com/google/jax/issues/5370
+    scipy_fun = lambda a, b: osp_special.logsumexp(a, b=b)
+    lax_fun = lambda a, b: lsp_special.logsumexp(a, b=b)
+    args_maker = lambda: [np.array([-1000, -2]), np.array([1, 0])]
+    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker)
+    self._CompileAndCheck(lax_fun, args_maker)
+
   @parameterized.named_parameters(itertools.chain.from_iterable(
     jtu.cases_from_list(
         {"testcase_name": jtu.format_test_name_suffix(
Repository: google/jax

Issue:
Bug in jax.random.poisson for lam=0
Dear people from JAX,

Thank you for writing this software, I love it! I am using JAX for some optics simulations for which I need to use the poissonian random number generator to add noise to my simulated measurements. I noticed that `poisson(key, lam=0)` always returns -1, which makes no sense, and it does not match with the numpy implementation.

In JAX, I would write:

```python
import jax.numpy as jnp
import jax
key = jax.random.PRNGKey(0)
original = jnp.zeros(10,)
x = jax.random.poisson(key,original,(10,))
print(x)
# [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1]
```
Numpy:

```python
import numpy as np
rng = np.random.default_rng()
s = rng.poisson(0, 10)
print(s)
# [0 0 0 0 0 0 0 0 0 0]
```

Looking at the source the problem seems to be in _poisson_knuth. An obvious solution would be to add a statement that just returns zero wherever lam==0,  but there's probably a more elegant solution?


Discussion:
Thanks for the report! I'll take a look.

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1173,11 +1173,12 @@ def _poisson(key, lam, shape, dtype):
   # λ -> ∞, so pick some arbitrary large value.
   lam_rejection = lax.select(use_knuth, lax.full_like(lam, 1e5), lam)
   max_iters = dtype.type(jnp.iinfo(dtype).max)  # insanely conservative
-  return lax.select(
-      use_knuth,
-      _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
-      _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
+  result = lax.select(
+    use_knuth,
+    _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
+    _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
   )
+  return lax.select(lam == 0, jnp.zeros_like(result), result)
 
 
 def poisson(key, lam, shape=(), dtype=dtypes.int_):


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -538,6 +538,12 @@ def testPoissonShape(self):
     x = random.poisson(key, np.array([2.0, 20.0]), shape=(3, 2))
     assert x.shape == (3, 2)
 
+  def testPoissonZeros(self):
+    key = random.PRNGKey(0)
+    lam = jnp.concatenate([jnp.zeros(10), 20 * jnp.ones(10)])
+    samples = random.poisson(key, lam, shape=(2, 20))
+    self.assertArraysEqual(samples[:, :10], jnp.zeros_like(samples[:, :10]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
       for dtype in jtu.dtypes.floating))
Repository: google/jax

Issue:
no constant handler for IntEnum
```python
import enum
import jax

class Foo(enum.IntEnum):
  bar = 1
  baz = 2


@jax.pmap
def f(_):
  return Foo.bar

f(jax.numpy.arange(1))
```

The same works with `jax.jit` instead of `jax.pmap`.


Discussion:

Resolving patch:
diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -29,11 +29,13 @@
 logging._warn_preinit_stderr = 0
 
 from ..config import flags
-from jax._src import util
+from jax._src import util, traceback_util
 from .. import dtypes
 import numpy as np
 import threading
 
+traceback_util.register_exclusion(__file__)
+
 try:
   from . import tpu_client
 except ImportError:
@@ -335,11 +337,12 @@ def constant(builder, py_val, canonicalize_types=True):
   Returns:
     A representation of the constant, either a ComputationDataHandle or None
   """
-  py_type = type(py_val)
-  if py_type in _constant_handlers:
-    return _constant_handlers[py_type](builder, py_val, canonicalize_types)
-  else:
-    raise TypeError("No constant handler for type: {}".format(py_type))
+  for t in type(py_val).mro():
+    handler = _constant_handlers.get(t)
+    if handler: return handler(builder, py_val, canonicalize_types)
+  if hasattr(py_val, '__jax_array__'):
+    return constant(builder, py_val.__jax_array__(), canonicalize_types)
+  raise TypeError("No constant handler for type: {}".format(type(py_val)))
 
 # HLO instructions optionally can be annotated to say how the output should be
 # spatially partitioned (represented in XLA as OpSharding protos, see


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -16,6 +16,7 @@
 import collections
 from contextlib import contextmanager
 import copy
+import enum
 from functools import partial
 import re
 import unittest
@@ -2359,6 +2360,19 @@ def __jax_array__(self):
     for f in [jnp.isscalar, jnp.size, jnp.shape, jnp.dtype]:
       self.assertEqual(f(x), f(a))
 
+  def test_constant_handler_mro(self):
+    # https://github.com/google/jax/issues/6129
+
+    class Foo(enum.IntEnum):
+      bar = 1
+
+    @api.pmap
+    def f(_):
+      return Foo.bar
+
+    ans = f(jnp.arange(1))  # doesn't crash
+    expected = jnp.arange(1) + 1
+    self.assertAllClose(ans, expected)
 
 
 class RematTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
betainc batching rule is not quite correct
This issue can be replicated using the following code

```python
import jax
import jax.numpy as jnp
from jax.scipy.special import betainc

def f(x):
    return betainc(jnp.ones(3), 1., x)

assert f(jnp.ones(3)).shape == (3,)
assert jax.vmap(f)(jnp.ones((2, 3))).shape == (2, 3)
# TypeError: regularized_incomplete_beta got arrays of different rank: (2, 3), (2,), (2, 3).
```


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2113,6 +2113,7 @@ def _naryop_weak_type_rule(name, *avals, **kwargs):
   return all(aval.weak_type for aval in avals)
 
 def naryop(result_dtype, accepted_dtypes, name, translation_rule=None):
+  # TODO(frostig,mattjj): only used with arity > 2 once, simplify
   dtype_rule = partial(naryop_dtype_rule, result_dtype, accepted_dtypes, name)
   shape_rule = partial(_broadcasting_shape_rule, name)
   weak_type_rule = partial(_naryop_weak_type_rule, name)
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -331,7 +331,8 @@ def broadcast_batcher(prim, args, dims, **params):
     return (out, (d,) * len(out)) if prim.multiple_results else (out, d)
   else:
     size, = {shape[d] for shape, d in shapes if d is not not_mapped}
-    args = [bdim_at_front(x, d, size) for x, d in zip(args, dims)]
+    args = [bdim_at_front(x, d, size) if np.ndim(x) else x
+            for x, d in zip(args, dims)]
     ndim = max(np.ndim(x) for x in args)  # special-case scalar broadcasting
     args = [_handle_scalar_broadcasting(ndim, x, d) for x, d in zip(args, dims)]
     out = prim.bind(*args, **params)


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -21,6 +21,7 @@
 
 import jax
 import jax.numpy as jnp
+import jax.scipy as jsp
 from jax import test_util as jtu
 from jax import lax
 from jax._src.lax import parallel
@@ -1240,5 +1241,13 @@ def testNonJaxTypedOutput(self):
       TypeError, "Output from batched function.*is not a valid JAX type"):
       vmap(lambda x: "hello")(np.arange(5))
 
+  def testIssue6096(self):
+    def f(x):
+      return jsp.special.betainc(jnp.ones(3), 1., x)
+
+    self.assertEqual(f(jnp.ones(3)).shape, (3,))
+    self.assertEqual(jax.vmap(f)(jnp.ones((2, 3))).shape, (2, 3))
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax

Issue:
NumPy compatibility of jax.numpy.linalg.eig for real-valued inputs
`jax.numpy.linalg.eig` for real-valued input always gives complex-valued output. However, NumPy casts the result to real dtype if the imaginary part of eigenvalues is all zeros.

```python
In [1]: import jax.numpy as jnp
In [2]: import numpy as np
In [3]: np.linalg.eig(np.diag((1, 2, 3)))
   ...: 
Out[3]: 
(array([1., 2., 3.]),
 array([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]]))
In [4]: jnp.linalg.eig(np.diag((1, 2, 3)))
   ...: 
Out[4]: 
[DeviceArray([1.+0.j, 2.+0.j, 3.+0.j], dtype=complex64),
 DeviceArray([[1.+0.j, 0.+0.j, 0.+0.j],
              [0.+0.j, 1.+0.j, 0.+0.j],
              [0.+0.j, 0.+0.j, 1.+0.j]], dtype=complex64)]
```

Is this the intended behavior?



Discussion:
Yes, this is intended. Due to JAX's XLA compilation model, it is not possible to make the return type of a function dependent on the values contained in the returned array.
There's probably still an action item here: we should document the difference.

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
@@ -253,7 +253,11 @@ def _det_jvp(primals, tangents):
   return y, jnp.trace(z, axis1=-1, axis2=-2)
 
 
-@_wraps(np.linalg.eig)
+@_wraps(np.linalg.eig, lax_description="""
+This differs from ``numpy.linalg.eig`` in that the return type of
+``jax.numpy.linalg.eig`` is always ``complex64`` for 32-bit input,
+and ``complex128`` for 64-bit input.
+""")
 def eig(a):
   a = _promote_arg_dtypes(jnp.asarray(a))
   return lax_linalg.eig(a, compute_left_eigenvectors=False)


Test patch: