jax-ml/jax · error · ValueError

Custom Partitioning rules must return Sharding.

Error message

Custom Partitioning rules must return Sharding.

What it means

In custom partitioning, the user-supplied infer_sharding_from_operands (or propagate_user_sharding) callback must return a jax Sharding instance. This ValueError is raised when the callback returns something else (e.g. a string, opsharding proto, or None).

Source

Thrown at jax/_src/custom_partitioning.py:147

    user_shapes = (shape,)
    user_shardings = (user_sharding,)
  user_shape = info.out_tree.unflatten(
      [
          info.unflatten_arg_shape(s, sharding)
          for s, sharding in zip(user_shapes, user_shardings)
      ]
  )
  result_sharding = info.propagate_user_sharding(
      *info.static_args, info.mesh, user_shape
  )
  result_shardings = _flatten_sharding(
      info.out_tree, result_sharding, user_shapes)
  return _pack_result_sharding(shape, result_shardings)


def _to_hlo_sharding(sharding, num_dimensions):
  if not isinstance(sharding, Sharding):
    raise ValueError("Custom Partitioning rules must return Sharding.")
  return sharding._to_xla_hlo_sharding(num_dimensions)


def _custom_partitioning_partition(arg_shapes, arg_shardings, result_shape,
                                   result_sharding, backend_string):
  info = _sharding_callbacks[backend_string]
  if result_shape.is_tuple():
    result_shapes = result_shape.tuple_shapes()
    result_shardings = result_sharding.tuple_elements()
  else:
    result_shapes = (result_shape,)
    result_shardings = (result_sharding,)
  mesh, lower_fn, result_sharding, arg_shardings = info.partition(
      *info.static_args,
      info.mesh,
      info.unflatten_arg_shapes(arg_shapes, arg_shardings),
      info.out_tree.unflatten(
          [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return a jax.sharding.Sharding subclass such as NamedSharding(mesh, P('x','y')) or GSPMDSharding
2. Ensure the callback's return annotation/type actually derives from jax.sharding.Sharding
3. Do not return HloSharding protos or strings; convert them first via GSPMDSharding if needed

Example fix

# before
def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
  return "{devices=[2,2] last_tile_dim_replicate}"  # string

# after
def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
  return NamedSharding(mesh, P('x', 'y'))

from jax.sharding import Sharding
out = infer_fn(mesh, arg_shapes, result_shape)
assert isinstance(out, Sharding), type(out)

def is_sharding(x):
    return isinstance(x, jax.sharding.Sharding)

When it happens

Trigger: A custom_partitioning function whose infer_sharding_from_operands callback returns e.g. an XLA HloSharding proto, a NamedSharding-like wrapper, or a plain tuple instead of a jax.sharding.Sharding subclass.

Common situations: Porting older TPU sharding code that returned op sharding protos; returning GSPMD sharding strings from callbacks.

Related errors

• Partitioned callback not implemented on {platform} backend.
• Mismatch in result shapes. %s vs %s
• Shardy is used, but sharding propagation callbacks instead o
• Sharding rule has {len(rule.operand_mappings)} operands, but
• Mosaic kernels cannot be automatically partitioned. Please w

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