{"record":{"id":"44a8348448967345","repo":"jax-ml/jax","slug":"custom-partitioning-rules-must-return-sharding","errorCode":null,"errorMessage":"Custom Partitioning rules must return Sharding.","messagePattern":"Custom Partitioning rules must return Sharding\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning.py","lineNumber":147,"sourceCode":"    user_shapes = (shape,)\n    user_shardings = (user_sharding,)\n  user_shape = info.out_tree.unflatten(\n      [\n          info.unflatten_arg_shape(s, sharding)\n          for s, sharding in zip(user_shapes, user_shardings)\n      ]\n  )\n  result_sharding = info.propagate_user_sharding(\n      *info.static_args, info.mesh, user_shape\n  )\n  result_shardings = _flatten_sharding(\n      info.out_tree, result_sharding, user_shapes)\n  return _pack_result_sharding(shape, result_shardings)\n\n\ndef _to_hlo_sharding(sharding, num_dimensions):\n  if not isinstance(sharding, Sharding):\n    raise ValueError(\"Custom Partitioning rules must return Sharding.\")\n  return sharding._to_xla_hlo_sharding(num_dimensions)\n\n\ndef _custom_partitioning_partition(arg_shapes, arg_shardings, result_shape,\n                                   result_sharding, backend_string):\n  info = _sharding_callbacks[backend_string]\n  if result_shape.is_tuple():\n    result_shapes = result_shape.tuple_shapes()\n    result_shardings = result_sharding.tuple_elements()\n  else:\n    result_shapes = (result_shape,)\n    result_shardings = (result_sharding,)\n  mesh, lower_fn, result_sharding, arg_shardings = info.partition(\n      *info.static_args,\n      info.mesh,\n      info.unflatten_arg_shapes(arg_shapes, arg_shardings),\n      info.out_tree.unflatten(\n          [","sourceCodeStart":129,"sourceCodeEnd":165,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning.py#L129-L165","documentation":"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).","triggerScenarios":"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.","commonSituations":"Porting older TPU sharding code that returned op sharding protos; returning GSPMD sharding strings from callbacks.","solutions":["Return a jax.sharding.Sharding subclass such as NamedSharding(mesh, P('x','y')) or GSPMDSharding","Ensure the callback's return annotation/type actually derives from jax.sharding.Sharding","Do not return HloSharding protos or strings; convert them first via GSPMDSharding if needed"],"exampleFix":"# before\ndef infer_sharding_from_operands(mesh, arg_shapes, result_shape):\n  return \"{devices=[2,2] last_tile_dim_replicate}\"  # string\n\n# after\ndef infer_sharding_from_operands(mesh, arg_shapes, result_shape):\n  return NamedSharding(mesh, P('x', 'y'))","handlingStrategy":"type-guard","validationCode":"from jax.sharding import Sharding\nout = infer_fn(mesh, arg_shapes, result_shape)\nassert isinstance(out, Sharding), type(out)","typeGuard":"def is_sharding(x):\n    return isinstance(x, jax.sharding.Sharding)","tryCatchPattern":null,"preventionTips":["Always return NamedSharding/GSPMDSharding objects from sharding callbacks","Never return proto strings from callbacks"],"tags":["jax","custom-partitioning","sharding","tpu"],"backgroundTag":"invalid-sharding-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}