Repository: google/jax

Issue:
pmap static_broadcasted_argnums does not respect keyword arguments.
Tested on internal TPU donut.

This code 

```python
import jax

def f(a, b):
  return a + b

g = jax.pmap(f, static_broadcasted_argnums=(1,))
g(jnp.ones((8, 1)), b=1)
```

Crashse with 
```
      5   return a + b
      6 g = jax.pmap(f, static_broadcasted_argnums=(1,))
----> 7 g(jnp.ones((8, 1)), b=1)

3 frames
google3/third_party/py/jax/api.py in _get_axis_size(i, shape, axis)
    775     return shape[axis]
    776   except (IndexError, TypeError) as e:
--> 777     raise ValueError(f"vmap got arg {i} of rank {len(shape)} but axis to be mapped {axis}") from e
    778 
    779 def _mapped_axis_size(tree, vals, dims, name):

ValueError: vmap got arg 1 of rank 0 but axis to be mapped 0
```

However, changing the last line to 
```python
g(jnp.ones((8, 1)), 1) # No keyword args!
```
works as expected.


Discussion:
Thanks for the report! Just checking: how does this compare to how `jit` works with `static_argnums` together with keyword arguments? (I'm not sure if this is a bug or just a bad error message, and perhaps comparing to `jit` can determine that.)
I just checked against `jit`, and indeed I think I'd currently consider this an issue about better errors, though we could upgrade both `jit` and `pmap` to try harder to reconcile arguments passed via keywords to parameter positions.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -138,9 +138,9 @@ def f_jitted(*args, **kwargs):
     if _jit_is_disabled():
       return fun(*args, **kwargs)
     if static_argnums and max(static_argnums) >= len(args):
-      msg = ("Jitted function has static_argnums={} but was called with only {}"
+      msg = ("jitted function has static_argnums={} but was called with only {}"
              " positional arguments.")
-      raise TypeError(msg.format(static_argnums, len(args)))
+      raise ValueError(msg.format(static_argnums, len(args)))
     f = lu.wrap_init(fun)
     if static_argnums:
       dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
@@ -771,21 +771,23 @@ def batched_fun(*args):
 
   return batched_fun
 
-def _get_axis_size(i:int, shape: Tuple[int, ...], axis: int):
+def _get_axis_size(name: str, i:int, shape: Tuple[int, ...], axis: int):
   try:
     return shape[axis]
   except (IndexError, TypeError) as e:
-    raise ValueError(f"vmap got arg {i} of rank {len(shape)} but axis to be mapped {axis}") from e
+    raise ValueError(f"{name} got arg {i} of rank {len(shape)} "
+                     f"but axis to be mapped {axis}") from e
 
 def _mapped_axis_size(tree, vals, dims, name):
-  mapped_axis_sizes = {_get_axis_size(i, onp.shape(x), d) for i, (x, d) in enumerate(zip(vals, dims))
-                        if d is not None}
+  mapped_axis_sizes = {_get_axis_size(name, i, onp.shape(x), d)
+                       for i, (x, d) in enumerate(zip(vals, dims))
+                       if d is not None}
   try:
     size, = mapped_axis_sizes
     return size
   except ValueError as e:
     if not mapped_axis_sizes:
-      raise ValueError("{} must have at least one non-None in_axes".format(name)) from e
+      raise ValueError(f"{name} must have at least one non-None value in in_axes") from e
     msg = "{} got inconsistent sizes for array axes to be mapped:\n".format(name) + "{}"
     # we switch the error message based on whether args is a tuple of arrays,
     # in which case we can produce an error message based on argument indices,
@@ -1033,7 +1035,14 @@ def pmap(fun: Callable, axis_name: Optional[AxisName] = None, *, in_axes=0,
   def f_pmapped(*args, **kwargs):
     f = lu.wrap_init(fun)
     if static_broadcasted_argnums:
-      dyn_argnums = [i for i in range(len(args)) if i not in static_broadcasted_argnums]
+      if max(static_broadcasted_argnums) >= len(args):
+        msg = ("pmapped function has static_broadcasted_argnums={} but was "
+               "called with only {} positional argument{}. All static "
+               "broadcasted arguments must be passed positionally.")
+        raise ValueError(msg.format(static_broadcasted_argnums, len(args),
+                                    "s" if len(args) > 1 else ""))
+      dyn_argnums = [i for i in range(len(args))
+                     if i not in static_broadcasted_argnums]
       f, dyn_args = argnums_partial(f, dyn_argnums, args)
       if isinstance(in_axes, tuple):
         dyn_in_axes = tuple(in_axes[i] for i in dyn_argnums)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1080,7 +1080,7 @@ def h(a, b):
       api.vmap(lambda x: x, in_axes=(jnp.array([1., 2.]),))(jnp.array([1., 2.]))
 
     with self.assertRaisesRegex(
-        ValueError, "vmap must have at least one non-None in_axes"):
+        ValueError, "vmap must have at least one non-None value in in_axes"):
       # If the output is mapped, there must be a non-None in_axes
       api.vmap(lambda x: x, in_axes=None)(jnp.array([1., 2.]))
 
@@ -1098,7 +1098,6 @@ def h(a, b):
       # If the output is mapped, then there must be some out_axes specified
       api.vmap(lambda x: x, out_axes=None)(jnp.array([1., 2.]))
 
-
   def test_vmap_structured_in_axes(self):
 
     A, B, C, D = 2, 3, 4, 5
@@ -1655,6 +1654,18 @@ def func1(x):
           re.DOTALL)):
       api.jit(func1)(2.)
 
+  def test_pmap_static_kwarg_error_message(self):
+    # https://github.com/google/jax/issues/3007
+    def f(a, b):
+      return a + b
+
+    g = jax.pmap(f, static_broadcasted_argnums=(1,))
+
+    msg = (r"pmapped function has static_broadcasted_argnums=\(1,\) but was "
+           r"called with only 1 positional argument. All static broadcasted "
+           r"arguments must be passed positionally.")
+    with self.assertRaisesRegex(ValueError, msg):
+      g(jnp.ones((1, 1)), b=1)
 
 class JaxprTest(jtu.JaxTestCase):