jax-ml/jax · error · TypeError

Missing required keyword argument: 'in_sharding'

Error message

Missing required keyword argument: 'in_sharding'

What it means

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.

Source

Thrown at jax/_src/pjit.py:2477

    return reshard(out, _out_sharding)
  return decorator


def explicit_axes(f=None, /, *, axes: str | tuple[str, ...] | None = None,
                  in_sharding=None):
  kwargs = dict(axes=axes, in_sharding=in_sharding)
  if f is None:
    return lambda g: _explicit_axes(g, **kwargs)
  return _explicit_axes(f, **kwargs)

def _explicit_axes(fun, *, axes, in_sharding):
  @wraps(fun)
  def decorator(*args, **kwargs):
    if in_sharding is None:
      if "in_sharding" in kwargs:
        _in_sharding = kwargs.pop("in_sharding")
      else:
        raise TypeError("Missing required keyword argument: 'in_sharding'")
    else:
      _in_sharding = in_sharding
    mesh_info = _get_new_mesh(axes, mesh_lib.AxisType.Explicit, 'explicit_axes')
    if mesh_info is None:
      raise ValueError(
          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'
          ' into a mesh context when using `explicit_axes` API.')
    with mesh_lib.use_abstract_mesh(mesh_info.new):
      args = reshard(args, _in_sharding)
      out = fun(*args, **kwargs)
    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(
        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)
    return reshard(out, out_specs)
  return decorator

# -------------------- with_layout_constraint --------------------

def with_layout_constraint(x, layouts):

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
f = explicit_axes_decorated(fun)  # missing in_sharding
f(x)
// after
f = explicit_axes_decorated(fun, in_sharding=P('data'))
f(x)

Defensive patterns

Strategy: validation

Validate before calling

import inspect
# before decorating
def has_in_sharding(decorator_kwargs, call_kwargs):
    return 'in_sharding' in decorator_kwargs or 'in_sharding' in call_kwargs
assert has_in_sharding(kwargs, {}), 'in_sharding required'

Type guard

def is_sharding_like(s) -> bool:
    import jax
    return isinstance(s, (jax.sharding.Sharding, jax.sharding.PartitionSpec)) or s is None

Try / catch

try:
    f(x)
except TypeError as e:
    if 'in_sharding' in str(e):
        f(x, in_sharding=P('data'))
    else:
        raise

Prevention

Always pass in_sharding explicitly at decoration time
Write a thin wrapper that asserts required kwargs before calling

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f2d474144ab5c6e4. Report an issue: GitHub.