{"record":{"id":"1014bae129d8e26f","repo":"jax-ml/jax","slug":"shardy-is-used-but-sharding-propagation-callbacks","errorCode":null,"errorMessage":"Shardy is used, but sharding propagation callbacks instead of sharding_rule are provided. Need to provide sharding_rule to migrate to Shardy.","messagePattern":"Shardy is used, but sharding propagation callbacks instead of sharding_rule are provided\\. Need to provide sharding_rule to migrate to Shardy\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning.py","lineNumber":541,"sourceCode":"    else:\n      static_args = ()\n      f_, dyn_args = self.fun, args\n    args_flat, in_tree = tree_util.tracing_registry.flatten(dyn_args)\n    in_avals = [core.typeof(x) for x in args_flat]\n    mesh = mesh_lib.thread_resources.env.physical_mesh\n    with core.extend_axis_env_nd(mesh.shape.items()):\n      closed_call, out_avals = pe.trace_to_jaxpr(\n          f_, ft.pack((ft.treedef_args_to_ft(in_tree, in_avals), {})), debug)\n    assert not closed_call.consts\n\n    propagate_user_sharding = None\n    infer_sharding_from_operands = None\n    sharding_rule = None\n    if config.use_shardy_partitioner.value:\n      if (self.sharding_rule is None and\n          (self.propagate_user_sharding is not None or\n            self.infer_sharding_from_operands is not None)):\n        raise NotImplementedError(\n            \"Shardy is used, but sharding propagation callbacks instead of \"\n            \"sharding_rule are provided. Need to provide sharding_rule to \"\n            \"migrate to Shardy.\"\n        )\n      sharding_rule = self.sharding_rule\n    else:\n      propagate_user_sharding = self.propagate_user_sharding\n      infer_sharding_from_operands = self.infer_sharding_from_operands\n\n    out_flat = custom_partitioning_p.bind(\n        *args_flat,\n        call=closed_call,\n        partition=self.partition,\n        propagate_user_sharding=propagate_user_sharding,\n        infer_sharding_from_operands=infer_sharding_from_operands,\n        decode_shardings=self.decode_shardings,\n        sharding_rule=sharding_rule,\n        in_tree=in_tree,","sourceCodeStart":523,"sourceCodeEnd":559,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning.py#L523-L559","documentation":"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.","triggerScenarios":"Using @custom_partitioning with infer_sharding_from_operands or propagate_user_sharding callbacks but no sharding_rule, while config.use_shardy_partitioner is True.","commonSituations":"Upgrading JAX to a version where Shardy is default and running older TPU sharding code relying on GSPMD callbacks.","solutions":["Provide a sharding_rule (Einsum-like string or SdyShardingRule) instead of the callbacks","If immediate migration is not possible, temporarily set jax_use_shardy_partitioner=False to restore GSPMD behavior (deprecated path)","Consult the JAX/Shardy migration guide to translate your propagation callback into a rule"],"exampleFix":"# before\n@custom_partitioning\ndef f(x):\n  ...\nf.def_partition(propagate_user_sharding=..., infer_sharding_from_operands=...)\n\n# after\nf.def_partition(sharding_rule='i j -> i j')","handlingStrategy":"fallback","validationCode":"from jax._src import config\nif config.use_shardy_partitioner.value:\n    assert sharding_rule is not None or (propagate_user_sharding is None and infer_sharding_from_operands is None)","typeGuard":null,"tryCatchPattern":"try:\n    f.def_partition(infer_sharding_from_operands=cb)\nexcept NotImplementedError as e:\n    if 'sharding_rule' in str(e):\n        jax.config.update('jax_use_shardy_partitioner', False)  # temporary\n        f.def_partition(infer_sharding_from_operands=cb)","preventionTips":["Migrate custom partitioners to sharding_rule strings ahead of JAX upgrades","Track the Shardy migration notices in release notes"],"tags":["jax","shardy","custom-partitioning","migration","tpu"],"backgroundTag":"legacy-api-removed-in-upgrade","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}