jax-ml/jax · error · TypeError

keyword arguments could not be resolved to positions

Error message

keyword arguments could not be resolved to positions

What it means

jax.custom_partitioning (and custom API sharding callbacks) resolves keyword arguments to positional ones via inspect.signature. If after binding and applying defaults there remain unresolvable keyword arguments (e.g. **kwargs in the wrapped function's signature or unrecognized kwargs), TypeError is raised.

Source

Thrown at jax/_src/custom_partitioning.py:55

from jax._src import flattree as ft
from jax._src import mesh as mesh_lib
from jax._src import sharding_impls
from jax._src import tree_util
from jax._src import xla_bridge as xb
from jax._src.custom_partitioning_sharding_rule import sdy_sharding_rule_to_mlir, SdyShardingRule, str_to_sdy_sharding_rule
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.lib import xla_client as xc
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import hlo
from jax._src.sharding import Sharding


def _resolve_kwargs(fun, args, kwargs):
  ba = inspect.signature(fun).bind(*args, **kwargs)
  ba.apply_defaults()
  if ba.kwargs:
    raise TypeError("keyword arguments could not be resolved to positions")
  else:
    return ba.args


class _ShardingCallbackInfo:

  def __init__(self, propagate_user_sharding, partition, to_mesh_pspec_sharding,
      in_tree, out_tree, infer_sharding_from_operands, module_context, mesh,
      static_args):
    self.propagate_user_sharding = propagate_user_sharding
    self.partition = partition
    self.to_mesh_pspec_sharding = to_mesh_pspec_sharding
    self.in_tree = in_tree
    self.out_tree = out_tree
    self.infer_sharding_from_operands = infer_sharding_from_operands
    self.module_context = module_context
    self.mesh = mesh
    self.static_args = static_args

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove **kwargs from the wrapped function's signature or stop passing the extra keyword
2. Pass all arguments positionally
3. Fix the keyword name typo so it matches a declared parameter

Example fix

# before
def f(x, **kwargs): ...
named_sharding_constraint(f)  # f called with kwargs later

# after
def f(x, scale): ...
# call f(x, scale=2.0) or f(x, 2.0) positionally

import inspect
ba = inspect.signature(fun).bind(*args, **kwargs)
ba.apply_defaults()
assert not ba.kwargs, f'unresolvable kwargs: {ba.kwargs}'

When it happens

Trigger: Passing a keyword argument that inspect.signature(fun).bind leaves in ba.kwargs — typically because fun accepts **kwargs, so the keyword cannot be mapped to a named parameter position.

Common situations: Wrapping a custom_partitioning-decorated function that has a **kwargs catch-all, or passing an unexpected kwarg name (typo) to a custom-partitioned function.

Related errors

• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me
• Default value must be of type bool, got {default} of type {g

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