Add kwargs support to in_shardings argument of jax.jit.

Currently, we only support this case:

* If kwargs are specified, then all in_shardings should be specified as dict matching the kwargs. args and kwargs mixture is not allowed. Either everything are kwargs or args hence in_shardings is a dict or specified positionally.

Example:

```
@partial(jax.jit, in_shardings=dict(y=s2, x=s1))
def f(x, y):
  return x * 2, y * 2

f(x=arr, y=arr2)
```

Fixes #17400

PiperOrigin-RevId: 623018032

Author: yashk2810

jax/_src/api_util.py

@@ -187,7 +187,8 @@ def _validate_argnums(sig: inspect.Signature, argnums: tuple[int, ...], argnums_
   inspect.Parameter.KEYWORD_ONLY,
 )
 def _validate_argnames(
-    sig: inspect.Signature, argnames: tuple[str, ...], argnames_name: str
+    sig: inspect.Signature, argnames: tuple[str, ...], argnames_name: str,
+    err_on_invalid: bool = False
 ) -> None:
   """
   Validate that the argnames are sensible for a given function.
@@ -212,13 +213,15 @@ def _validate_argnames(
   # Check whether any kwargs are invalid due to position only
   invalid_argnames = invalid_kwargs & set(argnames)
   if invalid_argnames:
-    # raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} "
-    #                  f"in {argnames_name}. These are positional-only")
-    # TODO: 2022-08-20 or later: replace with error
-    warnings.warn(f"Jitted function has invalid argnames {invalid_argnames} "
-                  f"in {argnames_name}. These are positional-only. "
-                  "This warning will be replaced by an error after 2022-08-20 "
-                  "at the earliest.", SyntaxWarning)
+    if err_on_invalid:
+      raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} "
+                      f"in {argnames_name}. These are positional-only")
+    else:
+      # TODO: 2022-08-20 or later: replace with error
+      warnings.warn(f"Jitted function has invalid argnames {invalid_argnames} "
+                    f"in {argnames_name}. These are positional-only. "
+                    "This warning will be replaced by an error after 2022-08-20 "
+                    "at the earliest.", SyntaxWarning)
 
   # Takes any kwargs
   if var_kwargs:
@@ -227,13 +230,16 @@ def _validate_argnames(
   # Check that all argnames exist on function
   invalid_argnames = set(argnames) - valid_kwargs
   if invalid_argnames:
-    # TODO: 2022-08-20 or later: replace with error
-    # raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} "
-    #                  f"in {argnames_name}. Function does not take these args.")
-    warnings.warn(f"Jitted function has invalid argnames {invalid_argnames} "
-                  f"in {argnames_name}. Function does not take these args."
-                  "This warning will be replaced by an error after 2022-08-20 "
-                  "at the earliest.", SyntaxWarning)
+    if err_on_invalid:
+      raise ValueError(f"Jitted function has invalid argnames {invalid_argnames} "
+                      f"in {argnames_name}. Function does not take these args. "
+                      f"Valid kwargs are {valid_kwargs}")
+    else:
+      # TODO: 2022-08-20 or later: replace with error
+      warnings.warn(f"Jitted function has invalid argnames {invalid_argnames} "
+                    f"in {argnames_name}. Function does not take these args."
+                    "This warning will be replaced by an error after 2022-08-20 "
+                    "at the earliest.", SyntaxWarning)
 
 
 def argnums_partial(f, dyn_argnums, args, require_static_args_hashable=True):

jax/_src/pjit.py

@@ -510,6 +510,35 @@ def make_jit(fun: Callable, in_shardings: Any, out_shardings: Any,
   return _make_jit_wrapper(jit_info)
 
 
+def _check_in_shardings_kwargs_compatibility(
+    kws: bool, has_args: bool, sig, in_shardings_treedef,
+    user_specified_in_shardings: bool):
+  if not user_specified_in_shardings:
+    return
+  node_data = in_shardings_treedef.node_data()
+  if node_data is None:
+    return
+  type_, keys = node_data
+  if not kws and type_ is not dict:
+    return
+  # TODO(yashkatariya): Try to allow these cases as and when need arises.
+  if has_args and kws and type_ is dict:
+    raise ValueError(
+        'If in_shardings is a dict, the function should be called with kwargs'
+        ' only.')
+  if kws and type_ is not dict:
+    raise ValueError(
+        'If you are using kwargs, in_shardings needs to be passed as a dict'
+        f' corresponding to the kwargs. Got in_shardings type: {type_}')
+  if not kws and type_ is dict:
+    raise ValueError(
+        'in_shardings can only be an instance of dict if you have kwargs.'
+        ' Please pass in_shardings positionally if you are using args.')
+  assert type_ is dict
+  api_util._validate_argnames(sig, tuple(keys), 'in_shardings',
+                              err_on_invalid=True)
+
+
 def _infer_params(jit_info, args, kwargs):
   (fun, fun_sourceinfo, fun_signature, user_specified_in_shardings,
    in_shardings_treedef, in_shardings_leaves, out_shardings_treedef,
@@ -519,9 +548,9 @@ def _infer_params(jit_info, args, kwargs):
    abstracted_axes, _, use_resource_env) = jit_info
 
   have_kwargs = bool(kwargs)
-  if have_kwargs and user_specified_in_shardings:
-    raise ValueError(
-        "pjit does not support kwargs when in_shardings is specified.")
+  _check_in_shardings_kwargs_compatibility(
+      have_kwargs, bool(args), fun_signature, in_shardings_treedef,
+      user_specified_in_shardings)
 
   if use_resource_env:
     # We need to fetch the mesh from inside the wrapped function, because
@@ -538,7 +567,6 @@ def _infer_params(jit_info, args, kwargs):
     pjit_mesh = None
     jit_name = 'jit'
 
-
   axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
 
   dbg = debug_info(jit_name, fun_sourceinfo, fun_signature, args, kwargs,
@@ -998,21 +1026,36 @@ def _process_in_axis_resources(in_shardings_treedef, in_shardings_leaves,
                                in_layouts_treedef, in_layouts_leaves,
                                in_avals, in_tree, debug_info,
                                device_or_backend_set, kws):
+  in_tree_args, in_tree_kwargs = treedef_children(in_tree)
   if not kws:
-    in_tree, _ = treedef_children(in_tree)
+    in_tree = in_tree_args
 
   orig_in_shardings = tree_unflatten(in_shardings_treedef, in_shardings_leaves)
   # Only do this if original in_shardings are unspecified. If it is AUTO, go
   # via flatten_axis_resources.
   if is_unspecified(orig_in_shardings):
     in_shardings_flat = (orig_in_shardings,) * len(in_avals)
+  elif isinstance(orig_in_shardings, dict):
+    if in_shardings_treedef != in_tree_kwargs:
+      # TODO(yashkatariya): Improve the error message drastically.
+      raise ValueError(
+          'Pytree of in_shardings and kwargs should be equal. Got in_shardings'
+          f' pytree: {in_shardings_treedef}, kwargs pytree: {in_tree_kwargs}')
+    in_shardings_flat = in_shardings_leaves
   else:
     in_shardings_flat = flatten_axis_resources(
         "pjit in_shardings", in_tree, orig_in_shardings, tupled_args=True)
 
   in_layouts = tree_unflatten(in_layouts_treedef, in_layouts_leaves)
   if in_layouts is None:
     in_layouts_flat = (in_layouts,) * len(in_avals)
+  elif isinstance(in_layouts, dict):
+    if in_layouts_treedef != in_tree_kwargs:
+      # TODO(yashkatariya): Improve the error message drastically.
+      raise ValueError(
+          'Pytree of in_layouts and kwargs should be equal. Got in_layouts'
+          f' pytree: {in_layouts_treedef}, kwargs pytree: {in_tree_kwargs}')
+    in_layouts_flat = in_layouts_leaves
   else:
     in_layouts_flat = flatten_axis_resources(
         "pjit in_layouts", in_tree, in_layouts, tupled_args=True)

tests/pjit_test.py

@@ -2822,13 +2822,6 @@ def f(x, y, z):
     self.assertEqual(cache_info3.hits, cache_info2.hits)
     self.assertEqual(cache_info3.misses, cache_info2.misses + 1)
 
-  def test_pjit_kwargs_axis_resources_error(self):
-    with self.assertRaisesRegex(
-        ValueError,
-        "pjit does not support kwargs when in_shardings is specified."):
-      pjit(lambda x: x,
-           in_shardings=SingleDeviceSharding(jax.devices()[0]))(x=jnp.arange(8.))
-
   def test_pjit_keep_unused_true(self):
     @partial(pjit, keep_unused=True)
     def f(x, y, z, a, b, c):  # pylint: disable=unused-argument
@@ -4015,6 +4008,120 @@ def f(*args):
     inps = [arr, *[inp] * 2001]
     f(inps)  # doesn't crash
 
+  def test_in_shardings_kwargs(self):
+    mesh = jtu.create_global_mesh((2,), 'x')
+    s = NamedSharding(mesh, P('x'))
+    s1 = NamedSharding(mesh, P())
+    np_inp = np.arange(8)
+    arr = jax.device_put(np_inp, s1)
+
+    kw_in_s = dict(y=s1, x=s)
+    @partial(jax.jit, in_shardings=kw_in_s, out_shardings=(s, s1))
+    def f(x, y):
+      return x * 2, y * 2
+
+    out1, out2 = f(x=np_inp, y=arr)
+    self.assertArraysEqual(out1, np_inp * 2)
+    self.assertArraysEqual(out2, np_inp * 2)
+    self.assertEqual(out1.sharding, s)
+    self.assertEqual(out2.sharding, s1)
+
+    c_in_s = f.lower(x=np_inp, y=np_inp).compile().input_shardings
+    _, kw_shardings = c_in_s
+    self.assertDictEqual(kw_shardings, kw_in_s)
+
+  def test_kwargs_leaf_in_shardings(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=s)
+    def f(x, y):
+      return x @ y.T
+
+    out = f(x=np_inp, y=np_inp)
+    self.assertArraysEqual(out, np_inp @ np_inp.T)
+
+  def test_arg_in_shardings_kwarg_error(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=dict(x=s, y=s))
+    def f(x, y):
+      return x * 2, y * 2
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'in_shardings can only be an instance of dict if you have kwargs.'
+        ' Please pass in_shardings positionally if you are using args.'):
+      f(np_inp, np_inp)
+
+  def test_kwarg_in_shardings_positional_error(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=(s, s))
+    def f(x, y):
+      return x * 2, y * 2
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'If you are using kwargs, in_shardings needs to be passed as a dict'
+        ' corresponding to the kwargs'):
+      f(x=np_inp, y=np_inp)
+
+  def test_invalid_in_shardings_kwarg(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=dict(z=s))
+    def f(x):
+      return x * 2
+
+    with self.assertRaisesRegex(
+        ValueError,
+        "Jitted function has invalid argnames {'z'} in in_shardings. Function"
+        " does not take these args. Valid kwargs are {'x'}"):
+      f(x=np_inp)
+
+  def test_mix_args_kwargs_in_shardings(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=dict(x=s, y=s, z=s))
+    def f(x, y, z):
+      return x * 2, y * 2, z * 2
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'If in_shardings is a dict, the function should be called with kwargs'
+        ' only.'):
+      f(np_inp, y=np_inp, z=np_inp)
+
+  def test_kwargs_in_shardings_partial(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=dict(x=s))
+    def f(x, y, z):
+      return x * 2, y * 2, z * 2
+
+    with self.assertRaisesRegex(
+        ValueError, "Pytree of in_shardings and kwargs should be equal"):
+      f(x=np_inp, y=np_inp, z=np_inp)
+
+  def test_args_kwargs_in_shardings_mixture(self):
+    s = SingleDeviceSharding(jax.devices()[0])
+    np_inp = np.arange(16).reshape(8, 2)
+
+    @partial(jax.jit, in_shardings=(s, dict(y=s, z=s)))
+    def f(x, y, z):
+      return x * 2, y * 2, z * 2
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'If you are using kwargs, in_shardings needs to be passed as a dict'):
+      f(np_inp, y=np_inp, z=np_inp)
+
 
 class TempSharding(Sharding):