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

Issue:
lax.scan has error when input array is not a jax array
lax.scan raise an TypeError: 'JaxTuple' object does not support indexing incorrectly

A simple repro can be found here 
```
import collections

from jax import lax
import jax.numpy as np

Point = collections.namedtuple('Point',
                               'x y')

p0 = Point(x=np.array(1), y=np.array(2))


def plus_one(p, iter_idx):
  return Point(p.x+1, p.y+1), iter_idx 

point, _ = lax.scan(plus_one, p0, list(range(5)))
print(point)
```

or in this colab https://colab.research.google.com/drive/1_jmdeIjfxJwUElNpwEbBR1gnAQDe8WbP


Discussion:

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -464,9 +464,21 @@ def _promote_aval_rank(n, xs):
 
 def _leading_dim_size(xs):
   if isinstance(xs, core.JaxTuple):
-    return _leading_dim_size(xs[0])
+    sizes = set(map(_leading_dim_size, xs))
+    if len(sizes) == 1:
+      return sizes.pop()
+    elif len(sizes) > 1:
+      msg = "scan got inconsistent leading axis sizes: {}"
+      raise ValueError(msg.format(sizes))
+    else:
+      raise ValueError("scan found no leading axis to scan over")
   else:
-    return xs.shape[0]
+    shape = onp.shape(xs)
+    if shape:
+      return shape[0]
+    else:
+      msg = "scan got value with no leading axis to scan over: {}"
+      raise ValueError(msg.format(xs))
 
 def _empty_arrays(aval):
   assert isinstance(aval, core.AbstractValue)
@@ -556,7 +568,7 @@ def scan(f, init, xs):
   carry_aval_out, y_aval = pv_out
   if carry_aval != carry_aval_out:
     msg = ("scanned function carry output does not match carry input: "
-           "input carry is {} and output carry is {}")
+           "input carry is {} and output carry is {}.")
     raise TypeError(msg.format(carry_aval, carry_aval_out))
   lifted_jaxpr = pe._closure_convert_jaxpr(jaxpr)
   consts_aval, _ = _abstractify(core.pack(consts))


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -16,6 +16,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 from functools import partial
 
 from absl.testing import absltest
@@ -623,6 +624,18 @@ def loss(test_params):
 
     api.grad(loss)(0.25)  # doesn't crash
 
+  def testIssue744(self):
+    Point = collections.namedtuple('Point', ['x', 'y'])
+    p0 = Point(x=np.array(1), y=np.array(2))
+
+    def plus_one(p, iter_idx):
+      return Point(p.x+1, p.y+1), iter_idx
+
+    self.assertRaisesRegexp(
+        ValueError,
+        'scan got value with no leading axis to scan over.*',
+        lambda: lax.scan(plus_one, p0, list(range(5))))
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
fix higher-order rev-mode lax.index_take test failures 
With the new-and-more-thorough `check_grads`, we started getting some failures in higher-order tests of `lax.index_take`. See commented-out line in lax_test.py's testIndexTake.


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
@@ -662,6 +662,7 @@ def scatter(operand, scatter_indices, updates, dimension_numbers):
 
 def index_take(src, idxs, axes):
   indices = concatenate([reshape(i, [i.shape[0], 1]) for i in idxs], 1)
+  indices = indices % onp.array([src.shape[ax] for ax in axes])
   slice_sizes = list(src.shape)
   for ax in axes:
     slice_sizes[ax] = 1


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2038,12 +2038,9 @@ def args_maker():
       ]
       for rng in [jtu.rand_default()]))
   def testIndexTakeGrad(self, shape, dtype, idxs, axes, rng):
-    idxs = tuple(rng(e.shape, e.dtype) for e in idxs)
     src = rng(shape, dtype)
     index_take = lambda src: lax.index_take(src, idxs, axes)
-    check_grads(index_take, (src,), 2, ["fwd"], 1e-2, 1e-2, 1)
-    # TODO(mattjj): fix rev mode failures here!
-    # check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
+    check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_idxs={}_dnums={}_slice_sizes={}".format(
Repository: google/jax

Issue:
implement scatter batching rule when indices are batched and updates are not
Implement _scatter_batching_rule when `scatter_indices_bdim is not None and updates_bdim is None`.

See https://colab.corp.google.com/drive/1VBna662Pcz8zOL3SlzSJ6eETrN2x90FF#scrollTo=vLHx5BcHPykw (internal only) for the source code + context.

When input_embeddings and unc_target do not have a batch dimension,
`total_loss, grads = inner_loss_value_and_grad(net_params, input_embeddings, unc_target)` works fine.
Adding a batch dimension to input_embeddings and unc_target and changing the line to
`total_loss, grads = vmap(lambda args: inner_loss_value_and_grad(net_params, args[0], args[1]))([input_embeddings, unc_target])`
gives this NotImplementedError error:

"""
~/virtualenvs/_3/lib/python3.7/site-packages/jax/lax/lax.py in _scatter_batching_rule(***failed resolving arguments***)
   2842 
   2843   if scatter_indices_bdim is not None and updates_bdim is None:
-> 2844     raise NotImplementedError  # TODO(mattjj,phawkins)
   2845   elif scatter_indices_bdim is None and updates_bdim is not None:
   2846     updates = batching.move_dim_to_front(updates, updates_bdim)

NotImplementedError: 
"""

Full stacktrace https://paste.googleplex.com/4748794925154304.




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
@@ -2860,8 +2860,10 @@ def _scatter_batching_rule(
   operand_bdim = 0
 
   if scatter_indices_bdim is not None and updates_bdim is None:
-    raise NotImplementedError  # TODO(mattjj,phawkins)
-  elif scatter_indices_bdim is None and updates_bdim is not None:
+    updates = broadcast(updates, (size,))
+    updates_bdim = 0
+
+  if scatter_indices_bdim is None and updates_bdim is not None:
     updates = batching.move_dim_to_front(updates, updates_bdim)
     inserted_window_dims = tuple(onp.add(1, dimension_numbers.inserted_window_dims))
     update_window_dims = (0,) + tuple(onp.add(1, dimension_numbers.update_window_dims))


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -31,6 +31,7 @@
 from jax.core import unit
 from jax.interpreters import partial_eval as pe
 from jax.util import partial, curry
+import jax.ops
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -911,6 +912,11 @@ def testEmptyTuples(self):
     self.assertAllClose(result, onp.array([1, 2]), check_dtypes=False)
     self.assertEqual((), empty_tuple)
 
+  def testIndexAddBatchedIndexesOnly(self):
+    f = lambda x, idx, y: jax.ops.index_add(x, jax.ops.index[idx], y)
+    result = vmap(f, (None, 0, None))(onp.zeros((10,)), onp.arange(10,), 1.)
+    self.assertAllClose(result, onp.eye(10), check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
add lax.pswapaxes pxla translation
Needed by @thenerdstation.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -106,6 +106,9 @@ def ppermute(x, axis_name, perm):
 def pswapaxes(x, axis_name, axis):
   """Swap the pmapped axis ``axis_name`` with the unmapped axis ``axis``.
 
+  The mapped axis size must be equal to the size of the unmapped axis; that is,
+  we must have ``lax.psum(1, axis_name) == x.shape[axis]``.
+
   This function is similar to ``psplit`` except the pmapped axis of the input is
   placed at the position ``axis`` in the output.
 
@@ -121,6 +124,12 @@ def pswapaxes(x, axis_name, axis):
     where ``axis_size`` is the size of the mapped axis named ``axis_name`` in
     the input ``x``.
   """
+  # TODO(mattjj): enable this check when _serial_pmap works with psum(1, ax)
+  # axis_size = psum(1, axis_name)
+  # if axis_size != x.shape[axis]:
+  #   msg = ("pswapaxes requires the size of the mapped axis ``axis_name`` equal "
+  #          "``x.shape[axis]``, but they are {} and {} respectively.")
+  #   raise ValueError(msg.format(axis_size(axis_name), x.shape[axis]))
   return pswapaxes_p.bind(x, axis_name=axis_name, axis=axis)
 
 def psplit(x, axis_name, axis):
@@ -233,7 +242,11 @@ def _pswapaxes_serial_pmap_rule(vals, axes, axis):
   perm[axis] = axis_in
   return lax.transpose(x, perm), axis_in
 
+def _pswapaxes_translation_rule(c, xla_x, axis, replica_groups):
+  return c.AllToAll(xla_x, axis, axis, replica_groups)
+
 pswapaxes_p = standard_pmap_primitive('pswapaxes')
+pxla.parallel_translation_rules[pswapaxes_p] = _pswapaxes_translation_rule
 parallel.serial_pmap_primitive_rules[pswapaxes_p] = _pswapaxes_serial_pmap_rule
 
 


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -519,6 +519,16 @@ def testVmapOfPmapTuple(self):
     self.assertAllClose(expected_bz1, bz1, check_dtypes=False)
     self.assertAllClose(bz2, bz2, check_dtypes=False)
 
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testPswapaxes(self):
+    device_count = xla_bridge.device_count()
+    shape = (device_count, 3, device_count, 5)
+    x = onp.arange(prod(shape)).reshape(shape)
+
+    ans = pmap(lambda x: lax.pswapaxes(x, 'i', 1), axis_name='i')(x)
+    expected = onp.swapaxes(x, 0, 2)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
DeviceTuple constant handler doesn't work (and is untested)
From @fehiepsi in [this comment](https://github.com/google/jax/issues/804#issuecomment-500169961).

```python
def f(x):
    return jit(lax.fori_loop, static_argnums=(2,))(0, 10, lambda *args: args[1], x)

jit(f)(1.)
```

```
TypeError: No constant handler for type: <class 'jaxlib.xla_extension.XlaOp'>
```


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
@@ -470,9 +470,8 @@ def __eq__(self, other):
 canonicalize_dtype_handlers[DeviceTuple] = identity
 
 def _device_tuple_constant_handler(c, val, canonicalize_types=True):
-  py_val = pack(c.Constant(elt, canonicalize_types=canonicalize_types)
-                for elt in val)
-  return c.Constant(py_val)
+  const = partial(c.Constant, canonicalize_types=canonicalize_types)
+  return c.Tuple(*map(const, val))
 xb.register_constant_handler(DeviceTuple, _device_tuple_constant_handler)
 
 # TODO(mattjj): could jit-compile a computation here


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1325,6 +1325,14 @@ def testLongConstantHandling(self):
       self.skipTest("Test is Python 2 specific")
     self.assertTrue(api.jit(lambda x: lax.lt(x, long(10)))(long(3)))
 
+  def testIssue831(self):
+    # Tests the DeviceTuple constant handler
+    def f(x):
+      g = lambda *args: args[1]
+      return api.jit(lax.fori_loop, static_argnums=(2,))( 0, 10, g, x)
+
+    api.jit(f)(1.)  # doesn't crash
+
 
 class DeviceConstantTest(jtu.JaxTestCase):
   def _CheckDeviceConstant(self, make_const, expected):
Repository: google/jax

Issue:
np.arange ignores dtype
arange ignores the passed dtype and is always `int*`
```python
a = np.arange(4, dtype=np.complex64)
print(a.dtype)
assert a.dtype == np.complex64
```
Results in:
```
int64
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-12-5b101f76f4c3> in <module>()
      1 a = np.arange(4, dtype=np.complex64)
      2 print(a.dtype)
----> 3 assert a.dtype is np.complex64

AssertionError: 
```


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
@@ -1379,7 +1379,7 @@ def identity(n, dtype=None):
 
 @_wraps(onp.arange)
 def arange(*args, **kwargs):
-  dtype = kwargs.pop("dtype", None)
+  dtype = kwargs.get("dtype", None)
   if not args:
     raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
 


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
@@ -1493,7 +1493,6 @@ def testIssue453(self):
     expected = onp.reshape(a, (3, 2), order='F')
     self.assertAllClose(ans, expected, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}_dtype={}".format(
           op, {bool: "bool", int: "int", float: "float"}[dtype]),
@@ -1538,6 +1537,10 @@ def testArange(self):
     self.assertFalse(type(lnp.arange(77)) == type(onp.arange(77)))
     self.assertTrue(type(lnp.arange(77)) == type(lax.iota(onp.int32, 77)))
 
+  def testIssue830(self):
+    a = lnp.arange(4, dtype=lnp.complex64)
+    self.assertEqual(a.dtype, lnp.complex64)
+
   def testIssue728(self):
     assert lnp.allclose(lnp.eye(5000), onp.eye(5000))
     self.assertEqual(0, onp.sum(lnp.eye(1050) - onp.eye(1050)))
Repository: google/jax

Issue:
vmap fails possibly due to internal mishandling of scalar
```python
import jax
import jax.numpy as np

x = np.zeros([3,2], dtype=np.float32)
y = np.zeros([3], dtype=np.float32)
foo = jax.vmap(np.vdot, (1,None), 0)
foo(x,y)
```
Gives `TypeError: transpose permutation isn't a permutation of operand dimensions, got permutation (1, 0) for operand shape (2L,).`


Discussion:
Inside `dot_batch_rule`, the two highlighted lines:
```
1926        if rbd is None:
1927          assert lbd in (0, 1)
1928          if lbd == 0:
1929            return dot(lhs, rhs), 0
1930          else:
1931 ->         return dot(T(rhs), lhs), 1
1932 
1933        if lbd is None:
1934          assert rbd in (0, 1)
1935          if rbd == onp.ndim(rhs) - 1:
1936 ->         return dot(lhs, rhs), 1
1937          else:
1938            return dot(rhs, T(lhs)), 0
```
are wrong in cases with a 1D non-batched argument (where the output will also be 1D, so the output `batch_dim` should be 0). Will fix.

Resolving patch:
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -349,6 +349,9 @@ def _moveaxis(force_bcast, sz, dst, src, aval, x):
         x = broadcast(x, sz, force_broadcast=force_bcast)
         src = 0
         dst_ = dst % (aval.ndim + 1)
+      elif src >= aval.ndim:
+        raise ValueError(
+          "cannot move axis {} in {}-dimensional array".format(src, aval.ndim))
       if src == dst_:
         return x
       else:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1927,12 +1927,12 @@ def _dot_batch_rule(batched_args, batch_dims):
       if lbd == 0:
         return dot(lhs, rhs), 0
       else:
-        return dot(T(rhs), lhs), 1
+        return dot(T(rhs), lhs), onp.ndim(rhs) - 1
 
     if lbd is None:
       assert rbd in (0, 1)
       if rbd == onp.ndim(rhs) - 1:
-        return dot(lhs, rhs), 1
+        return dot(lhs, rhs), onp.ndim(lhs) - 1
       else:
         return dot(rhs, T(lhs)), 0
 


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -290,6 +290,14 @@ def testDot3(self):
     expected = onp.einsum('inj,jk->nik', xs, ys)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testDot4(self):
+    R = onp.random.RandomState(0).randn
+    xs = R(3, 2)
+    ys = R(3)
+    ans = vmap(np.dot, in_axes=(1, None))(xs, ys)
+    expected = onp.einsum('ij,i->j', xs, ys)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testPad(self):
     R = onp.random.RandomState(0).randn
Repository: google/jax

Issue:
Backward pass of gather needlessly allocates global memory
When using a gather-type operation in the forward pass for word embeddings, the backward pass pre-allocates a constant all-zero matrix that is the same size as the embedding matrix.

My current fix is:
```
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2692,7 +2692,9 @@ def _gather_transpose_rule(t, operand, s
   assert operand is None
   if t is ad_util.zero:
     return [ad_util.zero, ad_util.zero]
-  zeros = broadcast(_const(t, 0), operand_shape)
+  zeros = full(operand_shape, 0, dtype=t.dtype)
   scatter_dnums = ScatterDimensionNumbers(
     update_window_dims=dimension_numbers.offset_dims,
     inserted_window_dims=dimension_numbers.collapsed_slice_dims,
```


Discussion:
This looks right to me. Please turn into a PR!

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2692,7 +2692,7 @@ def _gather_transpose_rule(t, operand, start_indices, dimension_numbers,
   assert operand is None
   if t is ad_util.zero:
     return [ad_util.zero, ad_util.zero]
-  zeros = broadcast(_const(t, 0), operand_shape)
+  zeros = full(operand_shape, 0, dtype=t.dtype)
   scatter_dnums = ScatterDimensionNumbers(
     update_window_dims=dimension_numbers.offset_dims,
     inserted_window_dims=dimension_numbers.collapsed_slice_dims,


Test patch:
Repository: google/jax

Issue:
np.where performs broadcasting in op-by-op mode
The jax version of `np.where` needs to broadcast its arguments to the same shape. When one of the arguments is a constant, the broadcast operation isn't jitted.

Examples:
- `np.where(dropout_mask, x, 0)` creates an all-zero constant of the same shape as `x` and stores it in global memory. This potentially wastes hundreds of MB of global memory. 
- A similar problem occurs for the first argument (condition) of `np.where` if it has fewer dimensions than the other two arguments.

The first instance can be worked around by doing `np.where(condition, x, np.zeros_like(x))`, but working around the second is trickier without resorting to `lax.tie_in`.


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
@@ -739,7 +739,11 @@ def _maybe_numpy_1_13_isclose_behavior(a, out):
     return out
 
 
+# The `jit` on `where` exists to avoid materializing constants in cases like
+# `np.where(np.zeros(1000), 7, 4)`. In op-by-op mode, we don't want to
+# materialize the broadcast forms of scalar arguments.
 @_wraps(onp.where)
+@jit
 def where(condition, x=None, y=None):
   if x is None or y is None:
     raise ValueError("Must use the three-argument form of where().")
@@ -1465,14 +1469,20 @@ def tri(N, M=None, k=0, dtype=None):
 
 @_wraps(onp.tril)
 def tril(m, k=0):
-  mask = tri(*shape(m)[-2:], k=k, dtype=bool)
-  return where(mask, m, zeros_like(m))
+  m_shape = shape(m)
+  if len(m_shape) < 2:
+    raise ValueError("Argument to jax.numpy.tril must be at least 2D")
+  mask = tri(*m_shape[-2:], k=k, dtype=bool)
+  return lax.select(lax.broadcast(mask, m_shape[:-2]), m, zeros_like(m))
 
 
 @_wraps(onp.triu)
 def triu(m, k=0):
-  mask = tri(*shape(m)[-2:], k=k - 1, dtype=bool)
-  return where(mask, zeros_like(m), m)
+  m_shape = shape(m)
+  if len(m_shape) < 2:
+    raise ValueError("Argument to jax.numpy.triu must be at least 2D")
+  mask = tri(*m_shape[-2:], k=k - 1, dtype=bool)
+  return lax.select(lax.broadcast(mask, m_shape[:-2]), zeros_like(m), m)
 
 
 @_wraps(onp.trace)


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
@@ -759,7 +759,7 @@ def testTri(self, m, n, k, dtype, rng):
        "dtype": dtype, "shape": shape, "op": op, "k": k,
        "rng": jtu.rand_default()}
       for dtype in default_dtypes
-      for shape in [shape for shape in all_shapes if len(shape) >= 1]
+      for shape in [shape for shape in all_shapes if len(shape) >= 2]
       for op in ["tril", "triu"]
       for k in list(range(-3, 3))))
   def testTriLU(self, dtype, shape, op, k, rng):
Repository: google/jax

Issue:
Grad should raise an error for non-inexact types
Continuation of https://github.com/google/jax/issues/48 .

`grad` should report an error for integer/bool types (i.e., non-inexact types):

e.g., this example should report an error:

```
def square(x):
  return x**2

val = 3
dfn = grad(square)
print(dfn(val))
```


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -784,6 +784,14 @@ def fun(*tangents):
 
   return apply_jaxtree_fun(fun, io_tree, *py_args)
 
+def _check_inexact_input_vjp(x):
+  aval = core.get_aval(x)
+  if not onp.issubdtype(aval.dtype, onp.inexact):
+    msg = ("Primal inputs to reverse-mode differentiation must be of float "
+           "or complex type, got type {}")
+    raise TypeError(msg.format(aval.dtype.name))
+
+
 def vjp(fun, *primals, **kwargs):
   """Compute a (reverse-mode) vector-Jacobian product of `fun`.
 
@@ -823,6 +831,7 @@ def vjp(fun, *primals, **kwargs):
     fun = lu.wrap_init(fun)
   primals_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, primals))
   _check_args(primals_flat)
+  tree_map(_check_inexact_input_vjp, primals)
   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
   if not has_aux:
     out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -815,6 +815,13 @@ def f(a, b, c):
                     if eqn.bound_subjaxprs)
     self.assertEqual(len(subjaxpr.eqns), 1)
 
+  def test_grad_of_int_errors(self):
+    dfn = grad(lambda x: x ** 2)
+    jtu.check_raises_regexp(
+      lambda: dfn(3), TypeError,
+      "Primal inputs to reverse-mode differentiation must be of float or "
+      "complex type, got type int..")
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -796,7 +796,7 @@ def testGatherBatchedBoth(self, op_axis, idxs_axis, shape, dtype, idxs, dnums,
        "op_axis": op_axis, "idxs_axis": idxs_axis, "shape": shape, "dtype":
        dtype, "idxs": idxs, "dnums": dnums, "slice_sizes": slice_sizes,
        "rng": rng, "rng_idx": rng_idx}
-      for dtype in [onp.float32, onp.int32]
+      for dtype in [onp.float32]
       for op_axis, idxs_axis, shape, idxs, dnums, slice_sizes in [
           (0, 0, (2, 5), onp.array([[[0], [2]], [[1], [3]]]),
            lax.GatherDimensionNumbers(