{"record":{"id":"f2d474144ab5c6e4","repo":"jax-ml/jax","slug":"missing-required-keyword-argument-in-sharding","errorCode":null,"errorMessage":"Missing required keyword argument: 'in_sharding'","messagePattern":"Missing required keyword argument: 'in_sharding'","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pjit.py","lineNumber":2477,"sourceCode":"    return reshard(out, _out_sharding)\n  return decorator\n\n\ndef explicit_axes(f=None, /, *, axes: str | tuple[str, ...] | None = None,\n                  in_sharding=None):\n  kwargs = dict(axes=axes, in_sharding=in_sharding)\n  if f is None:\n    return lambda g: _explicit_axes(g, **kwargs)\n  return _explicit_axes(f, **kwargs)\n\ndef _explicit_axes(fun, *, axes, in_sharding):\n  @wraps(fun)\n  def decorator(*args, **kwargs):\n    if in_sharding is None:\n      if \"in_sharding\" in kwargs:\n        _in_sharding = kwargs.pop(\"in_sharding\")\n      else:\n        raise TypeError(\"Missing required keyword argument: 'in_sharding'\")\n    else:\n      _in_sharding = in_sharding\n    mesh_info = _get_new_mesh(axes, mesh_lib.AxisType.Explicit, 'explicit_axes')\n    if mesh_info is None:\n      raise ValueError(\n          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'\n          ' into a mesh context when using `explicit_axes` API.')\n    with mesh_lib.use_abstract_mesh(mesh_info.new):\n      args = reshard(args, _in_sharding)\n      out = fun(*args, **kwargs)\n    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(\n        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)\n    return reshard(out, out_specs)\n  return decorator\n\n# -------------------- with_layout_constraint --------------------\n\ndef with_layout_constraint(x, layouts):","sourceCodeStart":2459,"sourceCodeEnd":2495,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pjit.py#L2459-L2495","documentation":"Thrown by JAX's explicit-axes pjit-style decorator when `in_sharding` is neither passed as a keyword argument to the decorator nor supplied at call time via kwargs. The API requires an input sharding specification before it can reshard arguments. It mirrors the pattern of required keyword-only arguments in JAX transform decorators.","triggerScenarios":"Calling the decorated function (or the decorator itself, e.g. `jit_explicit(...)` / `pjit` with `explicit_axes`) without providing `in_sharding=...`, either at decoration time or as a kwarg at call time.","commonSituations":"Migrating from older pjit APIs where in_axis_resources was positional; refactoring code that previously relied on a default sharding; forgetting the parameter when copying example code that uses explicit mesh axes.","solutions":["Pass `in_sharding=...` (e.g. a NamedSharding or PartitionSpec) as a keyword argument to the decorator","Alternatively supply `in_sharding` as a keyword when calling the decorated function","Check the explicit-axes API signature to confirm the expected sharding type"],"exampleFix":"// before\nf = explicit_axes_decorated(fun)  # missing in_sharding\nf(x)\n// after\nf = explicit_axes_decorated(fun, in_sharding=P('data'))\nf(x)","handlingStrategy":"validation","validationCode":"import inspect\n# before decorating\ndef has_in_sharding(decorator_kwargs, call_kwargs):\n    return 'in_sharding' in decorator_kwargs or 'in_sharding' in call_kwargs\nassert has_in_sharding(kwargs, {}), 'in_sharding required'","typeGuard":"def is_sharding_like(s) -> bool:\n    import jax\n    return isinstance(s, (jax.sharding.Sharding, jax.sharding.PartitionSpec)) or s is None","tryCatchPattern":"try:\n    f(x)\nexcept TypeError as e:\n    if 'in_sharding' in str(e):\n        f(x, in_sharding=P('data'))\n    else:\n        raise","preventionTips":["Always pass in_sharding explicitly at decoration time","Write a thin wrapper that asserts required kwargs before calling"],"tags":["jax","sharding","missing-argument","pjit"],"backgroundTag":"missing-required-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}