Reverts a1c8207

PiperOrigin-RevId: 623045488

Author: yashk2810

jax/_src/api_util.py

@@ -187,8 +187,7 @@ 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,
-    err_on_invalid: bool = False
+    sig: inspect.Signature, argnames: tuple[str, ...], argnames_name: str
 ) -> None:
   """
   Validate that the argnames are sensible for a given function.
@@ -213,15 +212,13 @@ def _validate_argnames(
   # Check whether any kwargs are invalid due to position only
   invalid_argnames = invalid_kwargs & set(argnames)
   if invalid_argnames:
-    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)
+    # 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)
 
   # Takes any kwargs
   if var_kwargs:
@@ -230,16 +227,13 @@ def _validate_argnames(
   # Check that all argnames exist on function
   invalid_argnames = set(argnames) - valid_kwargs
   if invalid_argnames:
-    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)
+    # 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)
 
 
 def argnums_partial(f, dyn_argnums, args, require_static_args_hashable=True):

jax/_src/pjit.py

@@ -510,35 +510,6 @@ 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,
@@ -548,9 +519,9 @@ def _infer_params(jit_info, args, kwargs):
    abstracted_axes, _, use_resource_env) = jit_info
 
   have_kwargs = bool(kwargs)
-  _check_in_shardings_kwargs_compatibility(
-      have_kwargs, bool(args), fun_signature, in_shardings_treedef,
-      user_specified_in_shardings)
+  if have_kwargs and user_specified_in_shardings:
+    raise ValueError(
+        "pjit does not support kwargs when in_shardings is specified.")
 
   if use_resource_env:
     # We need to fetch the mesh from inside the wrapped function, because
@@ -567,6 +538,7 @@ 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,
@@ -1026,36 +998,21 @@ 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 = in_tree_args
+    in_tree, _ = treedef_children(in_tree)
 
   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,6 +2822,13 @@ 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
@@ -4008,120 +4015,6 @@ 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):