Key word arguments breaks pmap
This code
```python
def f(index1, index2):
  return np.arange(36).reshape(6, 6)[index1, index2]
g = jax.jit(jax.pmap(f))
g(index1=onp.asarray([1]), index2=onp.asarray([2]))
```

Causes this error.
```
google3/third_party/py/jax/api.py in _pmap_axis_size(args)
    724     return next(_axis_size(leaf) for arg in args for leaf in tree_leaves(arg))
    725   except StopIteration:
--> 726     raise ValueError("pmap requires a leading axis to map over.")
    727 
    728 def _axis_size(x):

ValueError: pmap requires a leading axis to map over.
```

Whereas the exact same code without the kwargs works correctly.

```python
 
def f(index1, index2):
  return np.arange(36).reshape(6, 6)[index1, index2]
g = jax.jit(jax.pmap(f))
g(onp.asarray([1]), onp.asarray([2]))
```
Which nicely prints
```
DeviceArray([8], dtype=int32)
```
as expected.
diff --git a/jax/api.py b/jax/api.py
index d8e96e258..8609a144d 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -711,7 +711,7 @@ def pmap(fun, axis_name=None):
 
   @wraps(fun)
   def f_pmapped(*args, **kwargs):
-    axis_size = _pmap_axis_size(args)
+    axis_size = _pmap_axis_size((args, kwargs))
     f = lu.wrap_init(fun)
     args_flat, in_tree = tree_flatten((args, kwargs))
     flat_fun, out_tree = flatten_fun_leafout(f, in_tree)
diff --git a/tests/batching_test.py b/tests/batching_test.py
index cf4aadeee..898ff5e52 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -955,6 +955,14 @@ class BatchingTest(jtu.JaxTestCase):
     result = vmap(f, (None, 0, None))(onp.zeros((10,)), onp.arange(10,), 1.)
     self.assertAllClose(result, onp.eye(10), check_dtypes=False)
 
+  def testIssue1170(self):
+    def f(index1, index2):
+      return np.arange(36).reshape(6, 6)[index1, index2]
+    g = jax.jit(jax.pmap(f))
+    ans = g(index1=onp.asarray([1]), index2=onp.asarray([2]))
+    expected = g(onp.asarray([1]), onp.asarray([2]))
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
 
 if __name__ == '__main__':
   absltest.main()
['tests/batching_test.py', 'jax/api.py']