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

Issue:
jax.grad(has_aux=True) leaks identity tracers
```
def f(x):
  aux = dict(identity=x, p1=x+1)
  return x ** 2, aux

jax.grad(f, has_aux=True)(3.)

(DeviceArray(6., dtype=float32),
 {'identity': Traced<ConcreteArray(3.0, weak_type=True):JaxprTrace(level=0/0)>,
  'p1': DeviceArray(4., dtype=float32)})
```


Discussion:
It's surprising that one of them is being unboxed properly...

A workaround for now is just to use `lax.stop_gradient(x)` when you build the dictionary.

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -79,6 +79,7 @@ def jvp_subtrace_aux(master, primals, tangents):
   aux_tracers = map(trace.full_raise, aux)
   out_primals, out_tangents = unzip2((t.primal, t.tangent) for t in ans_tracers)
   aux_primals, _            = unzip2((t.primal, t.tangent) for t in aux_tracers)
+  aux_primals = map(core.full_lower, aux_primals)
   yield (out_primals, out_tangents), aux_primals
 
 def linearize(traceable, *primals, **kwargs):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -34,7 +34,7 @@
 import jax
 import jax.numpy as np
 from jax import jit, grad, device_put, jacfwd, jacrev, hessian
-from jax import api, lax
+from jax import api, core, lax
 from jax.core import Primitive
 from jax.interpreters import ad
 from jax.interpreters import xla
@@ -495,6 +495,17 @@ def test_grad_and_aux_constant(self):
     self.assertEqual(g, grad(lambda x: x**3)(4.))
     self.assertEqual(aux, [4.**2, 4.])
 
+  def test_grad_and_aux_no_tracers(self):
+    # see https://github.com/google/jax/issues/1950
+    def f(x):
+      aux = dict(identity=x, p1=x+1)
+      return x ** 2, aux
+
+    _, aux = jax.grad(f, has_aux=True)(3.)
+    self.assertIsInstance(aux, dict)
+    for val in aux.values():
+      self.assertNotIsInstance(val, core.Tracer)
+
   def test_jvp_mismatched_arguments(self):
     self.assertRaisesRegex(
       TypeError,
Repository: google/jax

Issue:
`&` and `|` operators don't work for `float32` arguments
Example:
```python
a = jnp.array([0.0, 1.0])
b = a | a
```

results in `
RuntimeError: Invalid argument: Expected pred or integral type in argument to and/or operation; got F32.`


Discussion:
This should be an error (it's an error in classic NumPy too), but it should be caught earlier. Thanks for the report!

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1649,6 +1649,7 @@ def _brcast_to(x, shape):
 
 _num = _int | _float | _complex
 _any = _int | _float | _complex | _bool
+_bool_or_int = _int | _bool
 
 neg_p = standard_unop(_num, 'neg')
 ad.deflinear(neg_p, lambda t: [neg(t)])
@@ -1814,15 +1815,15 @@ def _pow_jvp_rhs(g, ans, x, y):
 ad.defjvp2(pow_p, _pow_jvp_lhs, _pow_jvp_rhs)
 _replace_zero = lambda x: select(eq(x, _const(x, 0)), _ones(x), x)
 
-not_p = standard_unop(_int | _bool, 'not')
+not_p = standard_unop(_bool_or_int, 'not')
 
-and_p = standard_naryop([_any, _any], 'and')
+and_p = standard_naryop([_bool_or_int, _bool_or_int], 'and')
 ad.defjvp_zero(and_p)
 
-or_p = standard_naryop([_any, _any], 'or')
+or_p = standard_naryop([_bool_or_int, _bool_or_int], 'or')
 ad.defjvp_zero(or_p)
 
-xor_p = standard_naryop([_any, _any], 'xor')
+xor_p = standard_naryop([_bool_or_int, _bool_or_int], 'xor')
 ad.defjvp_zero(xor_p)
 
 def _add_transpose(t, x, y):


Test patch:
Repository: google/jax

Issue:
build.py specifies bazel 0.24 but uses repo_env which is new in bazel 0.26
https://github.com/google/jax/blob/master/build/build.py#L310 specifies that the minimum allowed bazel version is 0.24, but https://github.com/google/jax/blob/master/build/build.py#L167-L169 uses `--repo_env` (as opposed to `--action_env`), which was added in bazel 0.26 (in https://github.com/bazelbuild/bazel/commit/d7702b1 ).


Discussion:
@hawkinsp as this appears to stem from https://github.com/google/jax/commit/3e9ce2f .
Oops! Sorry about that. The right fix is probably for us to bump the minimum version of bazel (0.24 is old at this point). I'll do that.

Unfortunately I think there are other issues preventing us from upgrading bazel, at least when using the full build_jaxlib_wheels.sh stack. We may wanna revert to using --action_env for now, although ideally we fix the issues with new bazels.
I may have been mistaken -- I was just able to build at least one jaxlib wheel with bazel 1.1.0. I'll try bumping the minimum version as suggested and see what happens when I try a fresh jaxlib build (I've been debugging a lot of build issues so it's possible I messed something else up along the way).

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -307,7 +307,7 @@ def main():
 
   # Find a working Bazel.
   bazel_path = get_bazel_path(args.bazel_path)
-  check_bazel_version(bazel_path, min_version="0.24.0", max_version=None)
+  check_bazel_version(bazel_path, min_version="0.26.0", max_version=None)
   print("Bazel binary path: {}".format(bazel_path))
 
   python_bin_path = get_python_bin_path(args.python_bin_path)


Test patch:
Repository: google/jax

Issue:
tensordot raises TypeError: exceptions must derive from BaseException
Setup:
```python
import jax.numpy as jnp
import numpy as np

shp_a = (3, 2, 2)
shp_c = (2,)

a = np.random.random(shp_a)
c = np.random.random(shp_c)
```
Obviously, `jax.numpy.tensordot(a, c, axis=2)` should not work, but it raises a wrong error (logically, this should be something like dimension mismatch or index error):
```console
$ jnp.tensordot(a, c)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-48-e1d9970c866f> in <module>
----> 1 jnp.tensordot(a, c)

/usr/lib/python3.8/site-packages/jax/numpy/lax_numpy.py in tensordot(a, b, axes, precision)
   2309     if axes > _min(a_ndim, b_ndim):
   2310       msg = "Number of tensordot axes (axes {}) exceeds input ranks ({} and {})"
-> 2311       raise msg.format(axes, a.shape, b.shape)
   2312     contracting_dims = tuple(range(a_ndim - axes, a_ndim)), tuple(range(axes))
   2313   elif type(axes) in (list, tuple) and len(axes) == 2:

TypeError: exceptions must derive from BaseException
```
while in NumPy
```console
$ np.tensordot(a, c)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-46-3f592bfd2cd1> in <module>
----> 1 np.tensordot(a, c)

<__array_function__ internals> in tensordot(*args, **kwargs)

/usr/lib/python3.8/site-packages/numpy/core/numeric.py in tensordot(a, b, axes)
   1070     else:
   1071         for k in range(na):
-> 1072             if as_[axes_a[k]] != bs[axes_b[k]]:
   1073                 equal = False
   1074                 break

IndexError: tuple index out of range

```


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2308,7 +2308,7 @@ def tensordot(a, b, axes=2, precision=None):
   if type(axes) is int:
     if axes > _min(a_ndim, b_ndim):
       msg = "Number of tensordot axes (axes {}) exceeds input ranks ({} and {})"
-      raise msg.format(axes, a.shape, b.shape)
+      raise TypeError(msg.format(axes, a.shape, b.shape))
     contracting_dims = tuple(range(a_ndim - axes, a_ndim)), tuple(range(axes))
   elif type(axes) in (list, tuple) and len(axes) == 2:
     ax1, ax2 = axes


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
@@ -814,6 +814,13 @@ def onp_fun(a, b):
                             tol=tol)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  def testTensordotErrors(self):
+    a = onp.random.random((3, 2, 2))
+    b = onp.random.random((2,))
+    self.assertRaisesRegex(
+      TypeError, "Number of tensordot axes.*exceeds input ranks.*",
+      lambda: lnp.tensordot(a, b, axes=2))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_{}".format(
           jtu.format_shape_dtype_string(lhs_shape, lhs_dtype),
Repository: google/jax

Issue:
Docstring for jax.lax.tie_in
`jax.lax.tie_in` seems to have subtle but important properties, but [has no docstring](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.tie_in.html#jax.lax.tie_in).  Could we add one?


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1043,6 +1043,17 @@ def top_k(operand, k):
   return top_k_p.bind(operand, k=k)
 
 def tie_in(x, y):
+  """Gives ``y`` a fake data dependence on ``x``.
+
+  When staging to XLA (e.g. running under jit or pmap), values that don't depend
+  on computation inputs are computed op-by-op, and folded into the XLA
+  computation as constants.
+
+  ``tie_in`` provides a way to explicitly stage values into the computation.
+  When staging to XLA and ``x`` is already staged, then the result of ``tie_in``
+  is ``y``, but staged to XLA. Downstream use of the result will also be staged
+  to XLA.
+  """
   return tie_in_p.bind(x, y)
 
 


Test patch:
Repository: google/jax

Issue:
Add pxla.ShardedDeviceArray to jaxval_zeros_likers
This is a one-line fix at https://github.com/google/jax/blob/master/jax/lax/lax.py#L1503

Just adding this as an issue as a reminder for @mattjj 


Discussion:
Just add `ad_util.jaxval_zeros_likers[pxla.ShardedDeviceArray] = zeros_like_array` below that line
A thought - does `pxla.ShardedDeviceArray` also need to be added to `jaxval_adders`?

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1498,9 +1498,10 @@ def zeros_like_array(x):
   return full_like(x, 0)
 
 for t in itertools.chain(dtypes.python_scalar_dtypes.keys(), array_types,
-                         [xla.DeviceArray]):
+                         [xla.DeviceArray, pxla.ShardedDeviceArray]):
   ad_util.jaxval_adders[t] = add
 ad_util.jaxval_zeros_likers[xla.DeviceArray] = zeros_like_array
+ad_util.jaxval_zeros_likers[pxla.ShardedDeviceArray] = zeros_like_array
 
 
 ### primitives


Test patch:
Repository: google/jax

Issue:
BatchNorm doesn't work with center or scale turned off
BatchNorm raises TypeError when using without center or scale or both. Seems pretty straightforward to fix (e.x. as a workaround I've set beta to None `beta = None if not center else beta[ed]` if center is False, the same for gamma but there might be some better solution).

Exception:
beta = beta[ed] # ed=(None, slice(None, None, None))
TypeError: tuple indices must be integers or slices, not tuple

Code:
```
from jax import np
from jax import random
from jax.experimental import stax


init_fn, apply_fn = stax.BatchNorm(axis=0, center=False, scale=False)
_, params = init_fn(random.PRNGKey(0), (-1, 1))

x = np.zeros((2, 1))
apply_fn(params, x)
```

Details:
jax==0.1.62
jaxlib==0.1.42


Discussion:

Resolving patch:
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -134,12 +134,10 @@ def apply_fun(params, x, **kwargs):
     # TODO(phawkins): np.expand_dims should accept an axis tuple.
     # (https://github.com/numpy/numpy/issues/12290)
     ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-    beta = beta[ed]
-    gamma = gamma[ed]
     z = normalize(x, axis, epsilon=epsilon)
-    if center and scale: return gamma * z + beta
-    if center: return z + beta
-    if scale: return gamma * z
+    if center and scale: return gamma[ed] * z + beta[ed]
+    if center: return z + beta[ed]
+    if scale: return gamma[ed] * z
     return z
   return init_fun, apply_fun
 


Test patch:
diff --git a/tests/stax_test.py b/tests/stax_test.py
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -212,6 +212,19 @@ def testIssue182(self):
     assert out_shape == out.shape
     assert onp.allclose(onp.sum(onp.asarray(out), -1), 1.)
 
+  def testBatchNormNoScaleOrCenter(self):
+    key = random.PRNGKey(0)
+    axes = (0, 1, 2)
+    init_fun, apply_fun = stax.BatchNorm(axis=axes, center=False, scale=False)
+    input_shape = (4, 5, 6, 7)
+    inputs = random_inputs(onp.random.RandomState(0), input_shape)
+
+    out_shape, params = init_fun(key, input_shape)
+    out = apply_fun(params, inputs)
+    means = onp.mean(out, axis=(0, 1, 2))
+    std_devs = onp.std(out, axis=(0, 1, 2))
+    assert onp.allclose(means, onp.zeros_like(means), atol=1e-4)
+    assert onp.allclose(std_devs, onp.ones_like(std_devs), atol=1e-4)
 
   def testBatchNormShapeNHWC(self):
     key = random.PRNGKey(0)
Repository: google/jax

Issue:
Print values inside a vmap function
How to print values with a vmapped function even without any jit?
From advi.py example, 

def elbo(logprob, rng, mean, log_std):
    **sample** = diag_gaussian_sample(rng, mean, log_std) 
    return logprob(sample) - diag_gaussian_logpdf(sample, mean, log_std)

def batch_elbo(logprob, rng, params, num_samples):
    rngs = random.split(rng, num_samples)
    vectorized_elbo = vmap(partial(elbo, logprob), in_axes=(0, None, None))
    return np.mean(vectorized_elbo(rngs, *params))

print(sample) in VS Code debug mode returns the below message
None
Traced<ShapedArray(float32[2])>with<BatchTrace(level=2/0)>

This seems to related to #1396. 
Even if I disabled all the jit, still I can't print values with the vmapp function.



Discussion:
Like `jit`, `vmap` also traces the function in order to find the operations that need to be vectorized.  The print statement will show the abstract values that JAX is using to trace the function, before vmap.

I do not think there is a way right now to print; we do not yet have the jax.print operation.
I got it. Hope to support jax.print soon.
Thank you

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -449,7 +449,18 @@ def __getattr__(self, name):
         return attr
 
   def __repr__(self):
-    return 'Traced<{}>with<{}>'.format(self.aval, self._trace)
+    base = pp('Traced<{}>with<{}>'.format(self.aval, self._trace))
+    contents = self._contents()
+    if contents:
+      base += pp('  with ') >> vcat(pp('{} = '.format(name)) >> pp_payload
+                                    for name, pp_payload in contents)
+    return str(base)
+
+  def _contents(self):
+    try:
+      return [(name, pp(repr(getattr(self, name)))) for name in self.__slots__]
+    except AttributeError:
+      return ()
 
   def __copy__(self):
     return self


Test patch:
Repository: google/jax

Issue:
`AssertionError` with new `custom_jvp` for `jax.nn.relu`
On this [commit](https://github.com/google/jax/commit/ca23be63fbaed20192cda5f921afe177ac8dcf4d) (currently just a few commits behind head, none of which look relevant), I'm finding that
```python
import jax
import jax.numpy as np
import collections

class F(collections.namedtuple("F", ["a"])):
  def __call__(self, x):
    return jax.nn.relu(self.a) * x

@jax.jit
def g(f, x):
  return f(x)

jax.grad(F(2.0))(0.)  # works, returns DeviceArray(2., dtype=float32)
jax.grad(g, argnums=(1,))(F(2.0), 0.)  # AssertionError
```
hits an `AssertionError` on this line: https://github.com/google/jax/blob/1bb9aaa88c09a90a115c0304c05b0fc25932523b/jax/interpreters/partial_eval.py#L239

Commenting out the `@jax.jit` on `g` or replacing `jax.nn.relu(self.a)` with `np.maximum(self.a, 0)` both avoid the error.


Discussion:

Resolving patch:
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -252,6 +252,7 @@ def _flatten_jvp(in_tree, *args):
   yield primals_out + tangents_out, out_tree
 
 def _custom_jvp_call_bind(prim, fun, jvp, *args):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)
@@ -490,6 +491,7 @@ def _flatten_bwd(in_tree, out_trees, *args):
   yield cts_in
 
 def _custom_vjp_call_bind(prim, fun, fwd, bwd, *args, out_trees):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2781,6 +2781,19 @@ def g(x):
     expected = jax.grad(f, 0)(2., 0.1) + jax.grad(f, 0)(2., 0.2)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_lowering_out_of_traces(self):
+    # https://github.com/google/jax/issues/2578
+
+    class F(collections.namedtuple("F", ["a"])):
+      def __call__(self, x):
+        return jax.nn.relu(self.a) * x
+
+    @jax.jit
+    def g(f, x):
+      return f(x)
+
+    jax.grad(g, argnums=(1,))(F(2.0), 0.)  # doesn't crash
+
 
 class DeprecatedCustomTransformsTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
jax_debug_nans does not work with complex numbers
Running with flag --jax_debug_nans:

With this this, I get no error (but I do get nans):
```
  @jax.jit
  def fn(x):
    return (x * 0.) / jnp.zeros_like(x)
  print(fn(1j * jnp.ones([2, 3])))
```

but with this, I get an error (as expected):
```
  @jax.jit
  def fn(x):
    return (x * 0.) / jnp.zeros_like(x)
  print(fn(jnp.ones([2, 3])))
```

Could we add support for nan detection on complex numbers?


Discussion:
The issue seems to be that the nan checker only looks for `floating` types, not `complexfloating`. Numpy can detect complex nans just fine, so this should be a one line change in this line:
https://github.com/google/jax/blob/master/jax/interpreters/xla.py#L282

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -279,7 +279,7 @@ def check_nans(prim, bufs):
 
 def _check_nans(name, xla_shape, buf):
   assert not xla_shape.is_tuple()
-  if dtypes.issubdtype(xla_shape.element_type(), onp.floating):
+  if dtypes.issubdtype(xla_shape.element_type(), onp.inexact):
     if onp.any(onp.isnan(buf.to_py())):
       msg = "invalid value (nan) encountered in {}"
       raise FloatingPointError(msg.format(name))


Test patch:
Repository: google/jax

Issue:
np.median returns nan on a list of all inf
I found the following discrepancy between numpy and jax:
```
In [1]: import numpy as np                                                                                                                                                                 

In [2]: np.median([np.inf, np.inf, np.inf])                                                                                                                                                
Out[2]: inf

In [3]: import jax.numpy as np                                                                                                                                                             

In [4]: np.median([np.inf, np.inf, np.inf])                                                                                                                                                
/Users/stephentu/anaconda3/lib/python3.7/site-packages/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Out[4]: DeviceArray(nan, dtype=float32)

In [5]: import jax                                                                                                                                                                         

In [6]: jax.__version__                                                                                                                                                                    
Out[6]: '0.1.59'
```

Seems like returning `np.inf` is the sensible thing to do here. 

@mattjj suggested the culprit might be here: 
https://github.com/google/jax/blob/75077a14414124c13decc65c4ab1c1ac74174b81/jax/numpy/lax_numpy.py#L3374-L3375


Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -3481,7 +3481,7 @@ def percentile(a, q, axis=None, out=None, overwrite_input=False,
 def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
     q = 0.5
     return quantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
-                    keepdims=keepdims)
+                    keepdims=keepdims, interpolation='midpoint')
 
 def _astype(arr, dtype):
   lax._check_user_dtype_supported(dtype, "astype")


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
@@ -2139,7 +2139,7 @@ def onp_fun(*args):
 
   @parameterized.named_parameters(jtu.cases_from_list(
         {"testcase_name":
-           "_op=median_a_shape={}_axis={}_keepdims={}".format(
+           "_a_shape={}_axis={}_keepdims={}".format(
              jtu.format_shape_dtype_string(a_shape, a_dtype),
              axis, keepdims),
          "a_rng": jtu.rand_default(),
Repository: google/jax

Issue:
Generator/yield-support
Functions with yield does not seem to be supported and gives funny results in jit.

```
from jax import jit

def yf(x):
    yield(x)
    
jyf = jit(yf)
print(yf, jyf)
print(next(yf(1.5)))
print(next(jyf(1.5)))
```
Gives:

> <function yf at 0x7f79f01a7e18> <function jit.<locals>.f_jitted at 0x7f79f01ab1e0>
1.5
Traced<ShapedArray(float32[], weak_type=True):JaxprTrace(level=-1/1)>


Discussion:
It is intentional that this is not supported.

I'm not sure whether we could detect this case and issue a more graceful error. If we could, we should do so. Otherwise I think the best we could do is add a note to the documentation of `jit`.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -26,6 +26,7 @@
 
 import collections
 import functools
+import inspect
 import itertools as it
 import threading
 from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
@@ -75,7 +76,9 @@
 
 def _check_callable(fun):
   if not callable(fun):
-    raise TypeError("Expected a callable value, got {}".format(fun))
+    raise TypeError(f"Expected a callable value, got {fun}")
+  if inspect.isgeneratorfunction(fun):
+    raise TypeError(f"Expected a function, got a generator function: {fun}")
 
 class _ThreadLocalState(threading.local):
   def __init__(self):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -877,7 +877,13 @@ def test_jit_device(self):
 
   def test_jit_of_noncallable(self):
     self.assertRaisesRegex(TypeError, "Expected a callable value.*",
-                            lambda: api.jit(3))
+                           lambda: api.jit(3))
+
+  def test_jit_of_generator(self):
+    def gen(x):
+      yield x
+    self.assertRaisesRegex(TypeError, "Expected a function, got a generator function.*",
+                           lambda: api.jit(gen))
 
   def test_issue_1062(self):
     # code from https://github.com/google/jax/issues/1062 @shoyer