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

Issue:
Inconsistent behaviour of arange
I have installed JAX on Colab as in the QuickStart notebook. Then I run this instruction:

```
import jax.numpy as np
print(np.arange(0.0, 1.0, 0.1))
```

And I get an array with a single value [0], which is different from:

```
import numpy as np
print(np.arange(0.0, 1.0, 0.1))
```

Where I (correctly) get an array with 11 values 0, 0.1, ..., 0.9, 1.0. Is this intended?


Discussion:
Not intended, that's a bug! A temporary workaround is to `import numpy as onp` and just use `onp.arange` here. 
Ok, thanks a lot! Love the library!

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
@@ -834,29 +834,17 @@ def eye(N, M=None, k=None, dtype=onp.dtype("float64")):
 
 @_wraps(onp.arange)
 def arange(*args, **kwargs):
-  nargs = len(args)
-  start, step = 0, 1
-  dtype = kwargs.pop("dtype", None)
-  if kwargs:
-    raise TypeError("arange only accepts 'dtype' kwarg, got {}".format(kwargs))
-  if nargs == 0:
-    raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
-  elif nargs == 1:
-    stop, = args
-    dtype = dtype or _dtype(stop)
-    return lax.iota(dtype, stop)  # avoids materializing
-  elif nargs == 2:
-    start, stop = args
-    dtype = dtype or onp.result_type(start, stop)
-  elif nargs == 3:
-    start, stop, step = args
-    dtype = dtype or onp.result_type(start, stop, step)
-  elif nargs == 4:
-    start, stop, step, dtype = args
-    dtype = dtype or onp.result_type(start, stop, step)
-
-  size = (stop - start - 1) // step + 1
-  return start + step * lax.iota(dtype, size)
+  # attempt to generate a lazy IotaConstant, otherwise fall back to raw numpy
+  # TODO(mattjj): add tests for this function, then re-enable
+  # dtype = kwargs.pop("dtype", None)
+  # if not args:
+  #   raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
+  # elif len(args) == 1 and not kwargs:
+  #   stop, = args
+  #   dtype = dtype or _dtype(stop)
+  #   if onp.issubdtype(dtype, onp.integer):
+  #     return lax.iota(dtype, stop)  # avoids materializing
+  return onp.arange(*args, **kwargs)
 
 
 @_wraps(onp.repeat)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 
 setup(
     name='jax',
-    version='0.1.12',
+    version='0.1.13',
     description='Differentiate, compile, and transform Numpy code.',
     author='JAX team',
     author_email='jax-dev@google.com',


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
@@ -845,6 +845,13 @@ def testOnpMean(self):
     ans = onp.mean(x)
     self.assertAllClose(ans, onp.array([1./3, 1./3, 1./3]), check_dtypes=False)
 
+  # TODO(mattjj): more exhaustive arange tests
+  def testArangeOnFloats(self):
+    # from https://github.com/google/jax/issues/145
+    expected = onp.arange(0.0, 1.0, 0.1)
+    ans = lnp.arange(0.0, 1.0, 0.1)
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
add Softmax layer to stax
test in examples/mnist_classifier.py,
```python
def accuracy(params, batch):    
  inputs, targets = batch    
  target_class = np.argmax(targets, axis=1)
  x = predict(params, inputs)
  print(x[0])
```
the final layer is LogSoftmax, but the output seems not correct, 
```python  
Starting training...
[-2.3113673 -2.6440005 -2.4797316 -1.79847   -1.6207608 -2.931935
 -3.303906  -2.7275395 -2.2099946 -2.1403143]
```


Discussion:
Can you say a bit more about what might be incorrect here? The sum of the elementwise-exp of those numbers is 1, which is intended.
I mean the sum of those numbers should be 1, not their elementwise-exp's.
That would be true of a softmax layer, but `LogSoftmax` intended to be a log-softmax layer.

Would it be useful to you for us to add a softmax layer (including the exp)?
sorry, i misunderstood log-softmax, maybe it will be more convenient if providing loss functions.
Got it! I'll add that to my todo list.

Resolving patch:
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -47,6 +47,11 @@ def logsoftmax(x, axis=-1):
   """Apply log softmax to an array of logits, log-normalizing along an axis."""
   return x - logsumexp(x, axis, keepdims=True)
 
+def softmax(x, axis=-1):
+  """Apply softmax to an array of logits, exponentiating and normalizing along an axis."""
+  unnormalized = np.exp(x - x.max(axis, keepdims=True))
+  return unnormalized / unnormalized.sum(axis, keepdims=True)
+
 def fastvar(x, axis, keepdims):
   """A fast but less numerically-stable variance calculation than np.var."""
   return np.mean(x**2, axis, keepdims=keepdims) - np.mean(x, axis, keepdims=keepdims)**2
@@ -146,7 +151,9 @@ def _elemwise_no_params(fun, **kwargs):
   return init_fun, apply_fun
 Tanh = _elemwise_no_params(np.tanh)
 Relu = _elemwise_no_params(relu)
+Exp = _elemwise_no_params(np.exp)
 LogSoftmax = _elemwise_no_params(logsoftmax, axis=-1)
+Softmax = _elemwise_no_params(softmax, axis=-1)
 Softplus = _elemwise_no_params(softplus)
 
 


Test patch:
diff --git a/tests/stax_test.py b/tests/stax_test.py
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -159,6 +159,17 @@ def testFanInConcat(self, input_shapes, axis):
     init_fun, apply_fun = stax.FanInConcat(axis)
     _CheckShapeAgreement(self, init_fun, apply_fun, input_shapes)
 
+  def testIsuse182(self):
+    init_fun, apply_fun = stax.Softmax
+    input_shape = (10, 3)
+    inputs = onp.arange(30.).astype("float32").reshape(input_shape)
+
+    out_shape, params = init_fun(input_shape)
+    out = apply_fun(params, inputs)
+
+    assert out_shape == out.shape
+    assert onp.allclose(onp.sum(onp.asarray(out), -1), 1.)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
Failing lax_numpy_indexing_test.py tests on Python 3.7
`python tests/lax_numpy_indexing_test.py --num_generated_cases=1` on Python 3.7 (fresh Anaconda install on a macOS). Numpy version 1.15.4.


Discussion:
For the record, I got this error in python3.7 using the jaxlib wheel we have on pypi:

```
(myenv) mattjj@mattjj:~/packages/jax$ python tests/lax_numpy_indexing_test.py --num_generated_cases=1
/usr/local/google/home/mattjj/miniconda2/envs/myenv/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: compiletime version 3.5 of module 'lapack' does not match runtime version 3.7
  return f(*args, **kwds)
Traceback (most recent call last):
  File "tests/lax_numpy_indexing_test.py", line 29, in <module>
    from jax import api
  File "/usr/local/google/home/mattjj/packages/jax/jax/__init__.py", line 18, in <module>
    import jax.numpy as np  # side-effecting import sets up operator overloads
  File "/usr/local/google/home/mattjj/packages/jax/jax/numpy/__init__.py", line 18, in <module>
    from . import linalg
  File "/usr/local/google/home/mattjj/packages/jax/jax/numpy/linalg.py", line 24, in <module>
    from .. import lax_linalg
  File "/usr/local/google/home/mattjj/packages/jax/jax/lax_linalg.py", line 32, in <module>
    from jaxlib import lapack
  File "jaxlib/lapack.pyx", line 72, in init lapack
    cdef char cdiag = 'U' if diag else 'N'
  File "jaxlib/lapack.pyx", line 37, in lapack.register_cpu_custom_call_target
    Shape = xla_client.Shape
  File "/usr/local/google/home/mattjj/miniconda2/envs/myenv/lib/python3.7/site-packages/jaxlib/xla_client.py", line 1561, in register_cpu_custom_call_target
    c_api.RegisterCpuCustomCallTarget(name, fn)
UnboundLocalError: local variable 'name' referenced before assignment
```

I just commented out all the code in lax_linalg.py to get this test to run.

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
@@ -1592,8 +1592,8 @@ def _rewriting_take(arr, idx, axis=0):
       axis += isinstance(elt, slice)   # advance axis index if not eliminated
     unexpanded_shape_itr = iter(result.shape)
     result_shape = tuple(1 if elt is None else next(unexpanded_shape_itr)
-                         for elt in canonical_idx if not isinstance(elt, int))
-    return lax.reshape(result, result_shape)
+                         for elt in canonical_idx if isinstance(elt, (type(None), slice)))
+    return lax.reshape(result, result_shape) if result_shape else result
 
   # Handle advanced indexing (non-tuple sequence, ndarray of dtype int or bool,
   # or a tuple with at least one sequence object).


Test patch:
Repository: google/jax

Issue:
jax.random.randint range must be valid
`random.randint` doesn't check that `minval` < `maxval`. If it's not, you get surprising results. For example,
`print(random.randint(random.PRNGKey(10003), (), 0, 0))`
yields -1234647498.

The best thing would be to throw an error message, but if that's not practical then it'd at least be nice to have the behavior be defined.


Discussion:
Great point.

Hmm... if the arguments are static constants then it's easy to check in Python, but if `minval` and/or `maxval` are computed in XLA then we don't have a runtime error system to hook into. Maybe in that case we could just have a less surprising behavior, like deterministically returning minval (by setting `maxval = max(minval, maxval)` or something).

Any other ideas?

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -245,6 +245,10 @@ def randint(key, shape, minval, maxval, dtype=onp.int32):
   if nbits not in (32, 64):
     raise TypeError("randint only accepts 32- or 64-bit dtypes.")
 
+  # if we don't have minval < maxval, just always return minval
+  # https://github.com/google/jax/issues/222
+  maxval = lax.max(lax.add(minval, onp.array(1, dtype)), maxval)
+
   # This algorithm is biased whenever (maxval - minval) is not a power of 2.
   # We generate double the number of random bits required by the dtype so as to
   # reduce that bias.


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -149,6 +149,10 @@ def testShuffle(self, dtype):
     self.assertFalse(onp.all(perm1 == x))  # seems unlikely!
     self.assertTrue(onp.all(onp.sort(perm1) == x))
 
+  def testIssue222(self):
+    x = random.randint(random.PRNGKey(10003), (), 0, 0)
+    assert x == 0
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
Support for numpy.take_along_axis
Requesting support for [this Numpy primitive](https://docs.scipy.org/doc/numpy-1.15.0/reference/generated/numpy.take_along_axis.html#numpy.take_along_axis).


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
@@ -1553,6 +1553,21 @@ def argsort(a, axis=-1, kind='quicksort', order=None):
     return perm
 
 
+@_wraps(onp.take_along_axis)
+def take_along_axis(arr, indices, axis):
+  if axis is None and ndim(arr) != 1:
+    return take_along_axis(arr.ravel(), indices.ravel(), 0)
+  elif ndim(arr) == 1:
+    return lax.index_take(arr, (indices,), (0,))
+  else:
+    all_indices = [lax.broadcasted_iota(_dtype(indices), shape(indices), i)
+                   for i in range(ndim(arr))]
+    all_indices[axis] = indices
+    all_indices = tuple(map(ravel, all_indices))
+    out_flat = lax.index_take(arr, all_indices, tuple(range(ndim(arr))))
+    return reshape(out_flat, shape(indices))
+
+
 ### Indexing
 
 


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
@@ -1006,6 +1006,25 @@ def testArgsortManually(self):
     expected = onp.argsort(x)
     self.assertAllClose(expected, ans, check_dtypes=False)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}_axis={}".format(
+          jtu.format_shape_dtype_string(shape, dtype), axis),
+       "rng": rng, "shape": shape, "dtype": dtype, "axis": axis}
+      for shape in [(3,), (3, 4), (3, 4, 5)]
+      for axis in itertools.chain(range(len(shape)), [-1], [None])
+      for dtype in default_dtypes
+      for rng in [jtu.rand_default()]))
+  def testTakeAlongAxis(self, shape, dtype, axis, rng):
+    def args_maker():
+      x = rng(shape, dtype)
+      i = onp.argsort(x, axis=axis)
+      return x, i
+
+    lnp_op = lambda x, i: lnp.take_along_axis(x, i, axis=axis)
+    onp_op = lambda x, i: onp.take_along_axis(x, i, axis=axis)
+    self._CheckAgainstNumpy(lnp_op, onp_op, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
NumPy 1.16.0 causes warnings
```jax/lax.py:769: FutureWarning: Conversion of the second argument of issubdtype from complex to np.complexfloating is deprecated. In future, it will be treated as np.complex128 == np.dtype(complex).type.```

(I don't have code to reproduce it right now, because any attempt to create a minimal example makes the warning disappear even though it doesn't make any sense.)


Discussion:
Thanks for catching this! That should be enough information for us to track it down.
For me even this simple example gives that error:
```
import jax.numpy as np
from jax import grad, jit

f = lambda xx: np.sin(xx)
gf = jit(grad(f))
print(gf(1.))
```
```
0.5403023
/global/u2/h/hejia/jax/jax/lax.py:749: FutureWarning: Conversion of the second argument of issubdtype from `complex` to `np.complexfloating` is deprecated. In future, it will be treated as `np.complex128 == np.dtype(complex).type`.
  if not any(onp.issubdtype(aval.dtype, t) for t in accepted_dtypes):
/global/u2/h/hejia/jax/jax/lib/xla_bridge.py:146: UserWarning: No GPU found, falling back to CPU.
  warnings.warn('No GPU found, falling back to CPU.')
```
I'm using a jupyter notebook. The "No GPU" warning disappears if I run the cell again, but the "complexfloating" warning is always there.

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -831,7 +831,7 @@ def _brcast_to(x, shape):
 
 _f32 = {onp.float32}
 _float = {onp.floating}
-_complex = {onp.complex}
+_complex = {onp.complexfloating}
 _complex_elem_types = {onp.float32, onp.float64}
 _int = {onp.integer}
 _bool = {onp.bool_}


Test patch:
Repository: google/jax

Issue:
jax missing scipy.special.expit
Would be possible to add gradients for `expit` and `logit`?


Discussion:

Resolving patch:
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -19,7 +19,7 @@
 import scipy.special as osp_special
 
 from .. import lax
-from ..numpy.lax_numpy import _wraps
+from ..numpy.lax_numpy import _wraps, asarray
 
 
 # need to create new functions because _wraps sets the __name__ attribute
@@ -28,3 +28,16 @@
 erf = _wraps(osp_special.erf)(lambda x: lax.erf(x))
 erfc = _wraps(osp_special.erfc)(lambda x: lax.erfc(x))
 erfinv = _wraps(osp_special.erfinv)(lambda x: lax.erf_inv(x))
+
+
+@_wraps(osp_special.logit)
+def logit(x):
+  x = asarray(x)
+  return lax.log(lax.div(x, lax.sub(lax._const(x, 1), x)))
+
+
+@_wraps(osp_special.expit)
+def expit(x):
+  x = asarray(x)
+  one = lax._const(x, 1)
+  return lax.div(one, lax.add(one, lax.exp(lax.neg(x))))


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
@@ -62,6 +62,8 @@ def op_record(name, nargs, dtypes, rng, diff_modes, test_name=None):
     op_record("erf", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
     op_record("erfc", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
     op_record("erfinv", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
+    op_record("logit", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
+    op_record("expit", 1, float_dtypes, jtu.rand_small_positive(), ["rev"]),
 ]
 
 CombosWithReplacement = itertools.combinations_with_replacement
Repository: google/jax

Issue:
bug in np.full / lax.full for DeviceArray wrapped scalars
If you pass in a jax-wrapped scalar as the fill-constant value to np.fill, it breaks, e.g.:
```python
np.full((1,1), np.array([1])[0])
```
this hits the condition in lax.full:
```python
 elif isinstance(fill_value, xla.DeviceValue):
    return FilledConstant(convert_element_type(fill_value, dtype), shape)
```
no matter what dtype is, the convert_element_type result is still a DeviceArray, then
this line in lax.FillConstant:
```python
class FilledConstant(xla.DeviceConstant):
  __slots__ = ["fill_value"]

  def __init__(self, fill_value, shape):
    assert type(fill_value) is onp.ndarray
```
wants the fill value to be an onp ndarray instead.

I'm not too certain about the type logic here to make a PR confidently.


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -462,7 +462,8 @@ def full(shape, fill_value, dtype):
   if onp.isscalar(fill_value) or type(fill_value) is onp.ndarray:
     return FilledConstant(onp.asarray(fill_value, dtype), shape)
   elif isinstance(fill_value, xla.DeviceValue):
-    return FilledConstant(convert_element_type(fill_value, dtype), shape)
+    val = onp.asarray(fill_value, dtype)
+    return FilledConstant(val, shape)
   else:
     return broadcast(convert_element_type(fill_value, dtype), shape)
 


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
@@ -1187,7 +1187,6 @@ def args_maker():
       self._CheckAgainstNumpy(lnp_op, onp_op, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_op, args_maker, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_n={}_increasing={}".format(
           jtu.format_shape_dtype_string([shape], dtype),
@@ -1207,6 +1206,10 @@ def testVander(self, shape, dtype, n, increasing, rng):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=False)
 
+  def testIssue330(self):
+    x = lnp.full((1, 1), lnp.array([1])[0])  # doesn't crash
+    self.assertEqual(x[0, 0], 1)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
add option to globally disable jit for better debugging
(Relates to #196)

It's harder to debug jitted functions than non-jitted functions, since it's essentially impossible to introspect their intermediate state. One option is to comment out the jit decorator, but then you have to remember to put it back when you're done debugging.

It'd be nice to have a mechanism (maybe a context manager?) to globally disable jitting for debugging purposes; once bugs are fixed, you don't have to reinsert all of your @jit statements.


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -26,8 +26,11 @@
 
 import itertools
 import operator as op
+import os
 
 import numpy as onp
+from contextlib import contextmanager
+from distutils.util import strtobool
 
 from . import core
 from . import linear_util as lu
@@ -47,10 +50,16 @@
 from .interpreters import ad
 from .interpreters import batching
 from .interpreters import parallel
+from .config import flags, config
 
 map = safe_map
 zip = safe_zip
 
+FLAGS = flags.FLAGS
+flags.DEFINE_bool("jax_disable_jit",
+                  strtobool(os.getenv("JAX_DISABLE_JIT", "False")),
+                  "Disable JIT compilation and just call original Python.")
+
 
 def jit(fun, static_argnums=()):
   """Sets up `fun` for just-in-time compilation with XLA.
@@ -71,6 +80,8 @@ def jit(fun, static_argnums=()):
   """
   @wraps(fun)
   def f_jitted(*args, **kwargs):
+    if _jit_is_disabled or config.read('jax_disable_jit'):
+      return fun(*args, **kwargs)
     f = lu.wrap_init(fun, kwargs)
     dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
     f, dyn_args = argnums_partial(f, dyn_argnums, args)
@@ -84,6 +95,15 @@ def f_jitted(*args, **kwargs):
   return f_jitted
 
 
+@contextmanager
+def disable_jit():
+  global _jit_is_disabled
+  _jit_is_disabled, prev_val = True, _jit_is_disabled
+  yield
+  _jit_is_disabled = prev_val
+_jit_is_disabled = False
+
+
 def grad(fun, argnums=0):
   """Creates a function which evaluates the gradient of `fun`.
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -321,6 +321,23 @@ def test_hessian_on_pytrees(self):
                 (onp.array([0., 0.]), onp.array([0., 2.])))
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_disable_jit(self):
+    effects = []
+
+    @api.jit
+    def f(x):
+      effects.append(1)
+      return x
+
+    with api.disable_jit():
+      f(2)
+      f(2)
+    assert len(effects) == 2
+
+    f(2)
+    f(2)
+    assert len(effects) == 3
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
jit inside while_loop causes an error
I'll help generate a repro tomorrow.


Discussion:

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -135,20 +135,23 @@ def write(v, node):
   env = {}
   consts_env = dict(zip(jaxpr.constvars, const_vals))
   write(core.unitvar, c.Tuple())
-  map(write, jaxpr.constvars, map(c.Constant, const_vals))
-  map(write, jaxpr.freevars, map(c.ParameterWithShape, freevar_shapes))
+  if const_vals:
+    map(write, jaxpr.constvars, map(c.Constant, const_vals))
+    map(write, jaxpr.freevars, map(c.ParameterWithShape, freevar_shapes))
+  else:
+    all_freevars = it.chain(jaxpr.constvars, jaxpr.freevars)
+    map(write, all_freevars, map(c.ParameterWithShape, freevar_shapes))
   map(write, jaxpr.invars, map(c.ParameterWithShape, arg_shapes))
   for eqn in jaxpr.eqns:
     in_nodes = map(read, eqn.invars)
     in_shapes = map(c.GetShape, in_nodes)
-    subcs = [jaxpr_computation(subjaxpr,
-                               [consts_env[b] for b in const_bindings],
-                               map(c.GetShape, map(read, freevar_bindings)),
+    subcs = [jaxpr_computation(subjaxpr, (),
+                               map(c.GetShape, map(read, const_bindings + freevar_bindings)),
                                *in_shapes)
              for subjaxpr, const_bindings, freevar_bindings
              in eqn.bound_subjaxprs]
-    subfuns = [(subc, tuple(map(read, freevar_bindings)))
-               for subc, (_, _, freevar_bindings)
+    subfuns = [(subc, tuple(map(read, const_bindings + freevar_bindings)))
+               for subc, (_, const_bindings, freevar_bindings)
                in zip(subcs, eqn.bound_subjaxprs)]
     ans = translation_rule(eqn.primitive)(c, *(subfuns + in_nodes), **eqn.params)
     out_nodes = xla_destructure(c, ans) if eqn.destructure else [ans]


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1142,6 +1142,38 @@ def loop_body(state):
     self.assertEqual(cloop(3, limit, 1), limit - 3)
     assert not effect[0]
 
+  def testWhileWithClosureJit(self):
+
+    def loop(init, local_limit, inc):
+
+      def loop_cond(state):
+        pos, _ = state
+        return lax.lt(pos, local_limit)
+
+      def loop_body(state):
+        effect[0] = True
+        pos, count = state
+        f = lambda pos, inc: (lax.add(pos, 1), lax.add(count, inc))
+        return api.jit(f)(pos, inc)
+
+      result = lax._while_loop(loop_cond, loop_body, (init, 0))
+      _, count = result
+      return count
+
+    cloop = api.jit(loop)
+
+    limit = 10
+    effect = [False]
+    self.assertEqual(loop(2, limit, 1), limit - 2)
+    assert effect[0]
+    effect[0] = False
+    self.assertEqual(cloop(2, limit, 1), limit - 2)
+    assert effect[0]
+    effect[0] = False
+    self.assertEqual(cloop(2, limit, 1), limit - 2)
+    self.assertEqual(cloop(3, limit, 1), limit - 3)
+    assert not effect[0]
+
   def testNestedWhileWithDynamicUpdateSlice(self):
     num = 5