{"record":{"id":"1f04373a84fe9f44","repo":"jax-ml/jax","slug":"custom-partitioned-function-function-r-does-not","errorCode":null,"errorMessage":"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).","messagePattern":"Custom-partitioned function (.+?) 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\\)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning.py","lineNumber":600,"sourceCode":"      raise AssertionError(\n          'Please file a bug at https://github.com/jax-ml/jax/issues')\n    am = axis_context.abstract_mesh\n    if am is not None:\n      mesh = mesh_lib.Mesh(np.array(devices).reshape(am.axis_sizes),\n                           am.axis_names)\n  elif isinstance(axis_context, sharding_impls.SPMDAxisContext):\n    devices = axis_context.mesh._flat_devices_tuple\n  else:\n    devices = None\n\n  if not devices or len(devices) == 1:\n    return mlir.lower_fun(\n        core.jaxpr_as_fun(call), multiple_results=True)(ctx, *values)\n\n  if (not config.use_shardy_partitioner.value and\n      infer_sharding_from_operands is None):\n    function = call.jaxpr.debug_info.func_src_info\n    raise NotImplementedError(\n        f\"Custom-partitioned function {function!r} does not support GSPMD \"\n        \"sharding propagation rules. GSPMD is deprecated; please upgrade \"\n        \"to and enable the Shardy partitioner \"\n        \"(jax_use_shardy_partitioner=True, which is the default).\"\n    )\n\n  def to_mesh_pspec_sharding(hlo_sharding: xc.HloSharding | None, ndim):\n    if hlo_sharding is None:\n      return hlo_sharding\n    if mesh.empty or not decode_shardings:\n      assert devices is not None\n      return sharding_impls.GSPMDSharding(devices, hlo_sharding)\n    pspec = sharding_impls.parse_flatten_op_sharding(\n        hlo_sharding, mesh)[0]\n    pspec = sharding_impls.PartitionSpec(*pspec, *((None,) * (ndim - len(pspec))))\n    return sharding_impls.NamedSharding(mesh, pspec)\n\n  sharding_callback_info = _ShardingCallbackInfo(propagate_user_sharding,","sourceCodeStart":582,"sourceCodeEnd":618,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning.py#L582-L618","documentation":"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.","triggerScenarios":"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.","commonSituations":"Explicitly disabling Shardy for compatibility while using new-style sharding-rule-only custom partitioners.","solutions":["Re-enable Shardy: set jax_use_shardy_partitioner=True (default) or remove the override","Provide an infer_sharding_from_operands callback for the GSPMD path","Ensure your custom_partitioning setup matches the partitioner backend selected by config"],"exampleFix":"# before\njax.config.update('jax_use_shardy_partitioner', False)\n# f only has sharding_rule -> error\n\n# after\njax.config.update('jax_use_shardy_partitioner', True)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't disable jax_use_shardy_partitioner unless your custom partitioner implements GSPMD callbacks","Add infer_sharding_from_operands if you must run under GSPMD"],"tags":["jax","gspmd","shardy","custom-partitioning","config"],"backgroundTag":"deprecated-feature-disabled","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}