jax-ml/jax · error · TypeError

shard_map requires a callable for its first argument, but go

Error message

shard_map requires a callable for its first argument, but got {f} of type {type(f)}.

What it means

jax.shard_map (and smap internally) requires its first argument to be a Python callable. This TypeError catches passing a non-callable such as a result array, a string, or the output of a previous computation instead of the function itself.

Source

Thrown at jax/_src/shard_map.py:254

    raise TypeError("smap out_axes must be an int, None, or (nested) container "
                    f"with those types as leaves, but got {out_axes}.")

  in_specs = (Infer if in_axes is Infer else
              tree_map(partial(_axes_to_pspec, axis_name), in_axes,
                       is_leaf=lambda x: x is None))
  out_specs = tree_map(partial(_axes_to_pspec, axis_name), out_axes,
                       is_leaf=lambda x: x is None)
  return _shard_map(f, mesh=None, in_specs=in_specs, out_specs=out_specs,
                    axis_names={axis_name}, check_vma=True, _smap=True)


@partial(traceback_util.api_boundary, repro_api_name="jax.shard_map")
def _shard_map[F: Callable](
    f: F, *, mesh: Mesh | AbstractMesh | None,
    in_specs: Specs, out_specs: Specs, axis_names: Set[AxisName],
    check_vma: bool, _smap: bool = False) -> F:
  if not callable(f):
    raise TypeError("shard_map requires a callable for its first argument, "
                    f"but got {f} of type {type(f)}.")

  @util.wraps(f)
  @traceback_util.api_boundary
  def wrapped(*args):
    nonlocal mesh, axis_names
    mesh, axis_names = _shmap_checks(
        mesh, axis_names, in_specs, out_specs, _smap)
    dbg = api_util.debug_info("shard_map", f, args, {})
    args_flat = ft.flatten(args)
    api_util.check_no_transformed_refs_args(lambda: dbg, args_flat)

    try:
      in_specs_flat = broadcast_prefix(
          in_specs, args, is_leaf=lambda x: x is None)
    except ValueError:
      e, *_ = prefix_errors(in_specs, args)
      raise e('shard_map in_specs') from None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the first positional argument is a function/lambda: shard_map(lambda x: x * 2, mesh=...)
2. Check that intermediate results weren't accidentally assigned to the name you pass

Example fix

// before
result = shard_map(computation(x), mesh=mesh, in_specs=P('i'), out_specs=P('i'))

// after
result = shard_map(computation, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

assert callable(f), f'shard_map first arg must be callable, got {type(f)}'

def is_callable(f) -> bool:
    return callable(f)

When it happens

Trigger: Calling shard_map(arr, mesh=...) instead of shard_map(f, mesh=...), or passing a lambda result, a module object, or a string expression.

Common situations: Missing parentheses elsewhere producing an array, refactoring where the function variable is shadowed, or copy-paste where a decorated result is passed instead of the function.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• lax.associative_scan: fn argument should be callable.
• shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.
• `axis_names` argument of shard_map should be of type `frozen
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li

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