jax-ml/jax · error · NotImplementedError

Shardy is used, but sharding propagation callbacks instead o

Error message

Shardy is used, but sharding propagation callbacks instead of sharding_rule are provided. Need to provide sharding_rule to migrate to Shardy.

What it means

When jax_use_shardy_partitioner is enabled (the default in recent JAX), custom_partitioning no longer accepts the legacy GSPMD propagation callbacks (propagate_user_sharding / infer_sharding_from_operands) without a sharding_rule. This NotImplementedError forces migration to Shardy sharding rules.

Source

Thrown at jax/_src/custom_partitioning.py:541

    else:
      static_args = ()
      f_, dyn_args = self.fun, args
    args_flat, in_tree = tree_util.tracing_registry.flatten(dyn_args)
    in_avals = [core.typeof(x) for x in args_flat]
    mesh = mesh_lib.thread_resources.env.physical_mesh
    with core.extend_axis_env_nd(mesh.shape.items()):
      closed_call, out_avals = pe.trace_to_jaxpr(
          f_, ft.pack((ft.treedef_args_to_ft(in_tree, in_avals), {})), debug)
    assert not closed_call.consts

    propagate_user_sharding = None
    infer_sharding_from_operands = None
    sharding_rule = None
    if config.use_shardy_partitioner.value:
      if (self.sharding_rule is None and
          (self.propagate_user_sharding is not None or
            self.infer_sharding_from_operands is not None)):
        raise NotImplementedError(
            "Shardy is used, but sharding propagation callbacks instead of "
            "sharding_rule are provided. Need to provide sharding_rule to "
            "migrate to Shardy."
        )
      sharding_rule = self.sharding_rule
    else:
      propagate_user_sharding = self.propagate_user_sharding
      infer_sharding_from_operands = self.infer_sharding_from_operands

    out_flat = custom_partitioning_p.bind(
        *args_flat,
        call=closed_call,
        partition=self.partition,
        propagate_user_sharding=propagate_user_sharding,
        infer_sharding_from_operands=infer_sharding_from_operands,
        decode_shardings=self.decode_shardings,
        sharding_rule=sharding_rule,
        in_tree=in_tree,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide a sharding_rule (Einsum-like string or SdyShardingRule) instead of the callbacks
2. If immediate migration is not possible, temporarily set jax_use_shardy_partitioner=False to restore GSPMD behavior (deprecated path)
3. Consult the JAX/Shardy migration guide to translate your propagation callback into a rule

Example fix

# before
@custom_partitioning
def f(x):
  ...
f.def_partition(propagate_user_sharding=..., infer_sharding_from_operands=...)

# after
f.def_partition(sharding_rule='i j -> i j')

from jax._src import config
if config.use_shardy_partitioner.value:
    assert sharding_rule is not None or (propagate_user_sharding is None and infer_sharding_from_operands is None)

try:
    f.def_partition(infer_sharding_from_operands=cb)
except NotImplementedError as e:
    if 'sharding_rule' in str(e):
        jax.config.update('jax_use_shardy_partitioner', False)  # temporary
        f.def_partition(infer_sharding_from_operands=cb)

When it happens

Trigger: Using @custom_partitioning with infer_sharding_from_operands or propagate_user_sharding callbacks but no sharding_rule, while config.use_shardy_partitioner is True.

Common situations: Upgrading JAX to a version where Shardy is default and running older TPU sharding code relying on GSPMD callbacks.

Related errors

• Custom Partitioning rules must return Sharding.
• Unknown keyword arguments: {sharding_rule_dict}
• Custom-partitioned function {function!r} does not support GS
• sharding_rule callable must produce either an SdyShardingRul
• reduce_axes argument to vjp is deprecated

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