{"record":{"id":"f00c156e2f7f6a73","repo":"jax-ml/jax","slug":"mismatch-in-result-shapes-s-vs-s","errorCode":null,"errorMessage":"Mismatch in result shapes. %s vs %s","messagePattern":"Mismatch in result shapes\\. (.+?) vs (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_partitioning.py","lineNumber":189,"sourceCode":"  module_context = info.module_context\n\n  result_shardings = _flatten_sharding(\n      info.out_tree, result_sharding, result_shapes)\n  arg_shardings = _flatten_sharding(info.in_tree, arg_shardings, arg_shapes)\n  tiled_args = [\n      _to_jax_shape(sharding.tile(s))\n      for sharding, s in zip(arg_shardings, arg_shapes)\n  ]\n  tiled_results = [\n      _to_jax_shape(sharding.tile(s))\n      for sharding, s in zip(result_shardings, result_shapes)\n  ]\n  closed_jaxpr = api.make_jaxpr(lower_fn, axis_env=list(mesh.shape.items()))(\n      *info.in_tree.unflatten(tiled_args)\n  )\n  if ([(o.shape, o.dtype) for o in closed_jaxpr.out_avals] !=\n      [(t.shape, t.dtype) for t in tiled_results]):\n    raise ValueError(\n        \"Mismatch in result shapes. %s vs %s\"\n        % (repr(closed_jaxpr.out_avals), repr(tiled_results))\n    )\n  axis_context = sharding_impls.SPMDAxisContext(mesh, frozenset(mesh.axis_names))\n  with core.extend_axis_env_nd(mesh.shape.items()):\n    module = mlir.build_mlir_module_helper(\n        closed_jaxpr,\n        name=\"tmp_xla_computation\",\n        platforms=module_context.platforms,\n        backend=module_context.backend,\n        axis_context=axis_context,\n    )\n  result_sharding = _pack_result_sharding(result_shape, result_shardings)\n  return mlir.module_to_bytecode(module), arg_shardings, result_sharding\n\n\ndef _custom_partitioning_infer_sharding_from_operands(arg_shapes, arg_shardings,\n                                                      result_shape,","sourceCodeStart":171,"sourceCodeEnd":207,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_partitioning.py#L171-L207","documentation":"During custom partitioning, JAX traces the user's partition function on tiled (per-shard) arguments and verifies the traced outputs match the tiled result shapes/dtypes declared by the sharding callbacks. A mismatch raises this ValueError.","triggerScenarios":"A custom_partitioning partition function that returns arrays with different shapes or dtypes than the original function's outputs (e.g. forgetting to untile results, returning per-device stacks, or changing dtype inside the partitioned computation).","commonSituations":"Writing a custom partitioner that manually slices with wrong shard sizes, transposes outputs, or returns float32 where the original returned bf16.","solutions":["Make the partition function's outputs exactly match the original function's output shapes and dtypes on the tiled inputs","Return dsr/untiled results in the same order and structure as the original outputs","Double-check result_shape from the sharding callback matches what your partition code actually computes"],"exampleFix":"# before\ndef partition_fn(mesh, *tiled_args):\n  return (tiled_args[0].reshape(-1),)  # wrong shard shape/dtype\n\n# after\ndef partition_fn(mesh, *tiled_args):\n  return (some_op(tiled_args[0]),)  # same (shape, dtype) as declared result","handlingStrategy":"validation","validationCode":"import numpy as np\nouts = partition_fn(mesh, *tiled_args)\nassert [(np.shape(o), o.dtype) for o in outs] == [(s.shape, s.dtype) for s in tiled_results]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Assert output shape/dtype equality in a CPU test of the partition function before running on TPU"],"tags":["jax","custom-partitioning","sharding","shape-mismatch"],"backgroundTag":"output-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}