jax-ml/jax · error · NotImplementedError

Custom-partitioned function {function!r} does not support GS

Error message

Custom-partitioned function {function!r} does not support GSPMD sharding propagation rules. GSPMD is deprecated; please upgrade to and enable the Shardy partitioner (jax_use_shardy_partitioner=True, which is the default).

What it means

Under the legacy (non-Shardy) partitioner, a custom-partitioned function must define infer_sharding_from_operands for the compiler to infer output shardings. If Shardy is disabled and no such callback exists, JAX raises NotImplementedError pointing to the deprecated GSPMD path.

Source

Thrown at jax/_src/custom_partitioning.py:600

      raise AssertionError(
          'Please file a bug at https://github.com/jax-ml/jax/issues')
    am = axis_context.abstract_mesh
    if am is not None:
      mesh = mesh_lib.Mesh(np.array(devices).reshape(am.axis_sizes),
                           am.axis_names)
  elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
    devices = axis_context.mesh._flat_devices_tuple
  else:
    devices = None

  if not devices or len(devices) == 1:
    return mlir.lower_fun(
        core.jaxpr_as_fun(call), multiple_results=True)(ctx, *values)

  if (not config.use_shardy_partitioner.value and
      infer_sharding_from_operands is None):
    function = call.jaxpr.debug_info.func_src_info
    raise NotImplementedError(
        f"Custom-partitioned function {function!r} does not support GSPMD "
        "sharding propagation rules. GSPMD is deprecated; please upgrade "
        "to and enable the Shardy partitioner "
        "(jax_use_shardy_partitioner=True, which is the default)."
    )

  def to_mesh_pspec_sharding(hlo_sharding: xc.HloSharding | None, ndim):
    if hlo_sharding is None:
      return hlo_sharding
    if mesh.empty or not decode_shardings:
      assert devices is not None
      return sharding_impls.GSPMDSharding(devices, hlo_sharding)
    pspec = sharding_impls.parse_flatten_op_sharding(
        hlo_sharding, mesh)[0]
    pspec = sharding_impls.PartitionSpec(*pspec, *((None,) * (ndim - len(pspec))))
    return sharding_impls.NamedSharding(mesh, pspec)

  sharding_callback_info = _ShardingCallbackInfo(propagate_user_sharding,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-enable Shardy: set jax_use_shardy_partitioner=True (default) or remove the override
2. Provide an infer_sharding_from_operands callback for the GSPMD path
3. Ensure your custom_partitioning setup matches the partitioner backend selected by config

Example fix

# before
jax.config.update('jax_use_shardy_partitioner', False)
# f only has sharding_rule -> error

# after
jax.config.update('jax_use_shardy_partitioner', True)

When it happens

Trigger: Setting jax_use_shardy_partitioner=False and calling a custom_partitioning function that only defines sharding_rule (or nothing) with no infer_sharding_from_operands.

Common situations: Explicitly disabling Shardy for compatibility while using new-style sharding-rule-only custom partitioners.

Related errors

• Unknown keyword arguments: {sharding_rule_dict}
• Shardy is used, but sharding propagation callbacks instead o
• sharding_rule callable must produce either an SdyShardingRul
• invalid truth value {val!r} for environment {varname!r}
• Unrecognized config option: {name}

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