jax-ml/jax · error · ValueError

Mismatch in result shapes. %s vs %s

Error message

Mismatch in result shapes. %s vs %s

What it means

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.

Source

Thrown at jax/_src/custom_partitioning.py:189

  module_context = info.module_context

  result_shardings = _flatten_sharding(
      info.out_tree, result_sharding, result_shapes)
  arg_shardings = _flatten_sharding(info.in_tree, arg_shardings, arg_shapes)
  tiled_args = [
      _to_jax_shape(sharding.tile(s))
      for sharding, s in zip(arg_shardings, arg_shapes)
  ]
  tiled_results = [
      _to_jax_shape(sharding.tile(s))
      for sharding, s in zip(result_shardings, result_shapes)
  ]
  closed_jaxpr = api.make_jaxpr(lower_fn, axis_env=list(mesh.shape.items()))(
      *info.in_tree.unflatten(tiled_args)
  )
  if ([(o.shape, o.dtype) for o in closed_jaxpr.out_avals] !=
      [(t.shape, t.dtype) for t in tiled_results]):
    raise ValueError(
        "Mismatch in result shapes. %s vs %s"
        % (repr(closed_jaxpr.out_avals), repr(tiled_results))
    )
  axis_context = sharding_impls.SPMDAxisContext(mesh, frozenset(mesh.axis_names))
  with core.extend_axis_env_nd(mesh.shape.items()):
    module = mlir.build_mlir_module_helper(
        closed_jaxpr,
        name="tmp_xla_computation",
        platforms=module_context.platforms,
        backend=module_context.backend,
        axis_context=axis_context,
    )
  result_sharding = _pack_result_sharding(result_shape, result_shardings)
  return mlir.module_to_bytecode(module), arg_shardings, result_sharding


def _custom_partitioning_infer_sharding_from_operands(arg_shapes, arg_shardings,
                                                      result_shape,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the partition function's outputs exactly match the original function's output shapes and dtypes on the tiled inputs
2. Return dsr/untiled results in the same order and structure as the original outputs
3. Double-check result_shape from the sharding callback matches what your partition code actually computes

Example fix

# before
def partition_fn(mesh, *tiled_args):
  return (tiled_args[0].reshape(-1),)  # wrong shard shape/dtype

# after
def partition_fn(mesh, *tiled_args):
  return (some_op(tiled_args[0]),)  # same (shape, dtype) as declared result

import numpy as np
outs = partition_fn(mesh, *tiled_args)
assert [(np.shape(o), o.dtype) for o in outs] == [(s.shape, s.dtype) for s in tiled_results]

When it happens

Trigger: 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).

Common situations: Writing a custom partitioner that manually slices with wrong shard sizes, transposes outputs, or returns float32 where the original returned bf16.

Related errors

• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• Custom Partitioning rules must return Sharding.
• Sharding rule has {len(rule.operand_mappings)} operands, but
• type of weights must match type of x. Got typeof(x)={core.ty
• Expected source shape to be {expected_src_shape}, but got {s

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