{"record":{"id":"6141ac6b6acf8c43","repo":"jax-ml/jax","slug":"array-inputs-to-associative-scan-must-have-the-sam","errorCode":null,"errorMessage":"Array inputs to associative_scan must have the same first dimension. (saw: {})","messagePattern":"Array inputs to associative_scan must have the same first dimension\\. \\(saw: (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2900,"sourceCode":"\n  def combine(a_flat, b_flat):\n    # Lower `fn` to operate on flattened sequences of elems.\n    a = tree_unflatten(tree, a_flat)\n    b = tree_unflatten(tree, b_flat)\n    c = fn(a, b)\n    c_flat, _ = tree_flatten(c)\n    return c_flat\n\n  # Check that all inputs have a consistent leading dimension `num_elems`.\n  axis = util.canonicalize_axis(axis, elems_flat[0].ndim)\n\n  if not core.is_constant_dim(elems_flat[0].shape[axis]):\n    raise NotImplementedError(\"associative scan over axis \"\n        f\"of non-constant size: {elems_flat[0].shape[axis]}. You may be \"\n        \"able to avoid this on TPU. See b/274176030.\")\n  num_elems = int(elems_flat[0].shape[axis])\n  if not all(int(elem.shape[axis]) == num_elems for elem in elems_flat[1:]):\n    raise ValueError('Array inputs to associative_scan must have the same '\n                     'first dimension. (saw: {})'\n                     .format([elem.shape for elem in elems_flat]))\n\n\n  # Summary of algorithm:\n  #\n  # Consider elements of `_scan(elems)` at odd indices. That's the same as first\n  # summing successive pairs of elements of `elems` and performing a scan on\n  # that half sized tensor. We perform the latter scan by recursion.\n  #\n  # Now consider the even elements of `_scan(elems)`. These can be computed\n  # from the odd elements of `_scan(elems)` by adding each odd element of\n  # `_scan(elems)` to the matching even element in the original `elems`.\n  #\n  # We return the odd and even elements interleaved.\n  #\n  # For the base case of the recursion we return the first element\n  # of `elems` followed by the sum of the first two elements computed as","sourceCodeStart":2882,"sourceCodeEnd":2918,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2882-L2918","documentation":"associative_scan flattens the (possibly pytree) input and requires every leaf array to have the same length along the scan axis so the combine tree is well-defined. After computing num_elems from the first leaf, it verifies all others match and raises ValueError listing all leaf shapes otherwise.","triggerScenarios":"Passing a pytree of arrays (tuple/dict) whose members have different lengths along `axis`, e.g. ((jnp.ones(5), jnp.ones(3)),) — note only the scanned axis must match, other dims may differ; or mixing arrays of different sequence lengths in a structured carry.","commonSituations":"Scanning over structured state (e.g. (cumsum, logprobs) tuples) where one element was sliced or padded differently; off-by-one slicing like x[:-1] applied to only one member of the tuple.","solutions":["Ensure every array in the pytree has identical size along the scan axis (pad or slice all members consistently)","If elements are independent, run separate associative_scan calls per array","Print [a.shape for a in jax.tree_util.tree_leaves(elems)] to find the mismatched leaf"],"exampleFix":"// before\nlax.associative_scan(fn, (jnp.ones(5), jnp.ones(4)))\n// after\nlax.associative_scan(fn, (jnp.ones(5), jnp.ones(5)))","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nleaves = jtu.tree_leaves(elems)\nn = leaves[0].shape[axis]\nassert all(l.shape[axis] == n for l in leaves), [l.shape for l in leaves]","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Tree-map slicing/padding so every leaf changes together","Unit-test scan pytrees with heterogeneous shapes"],"tags":["jax","associative-scan","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}