{"record":{"id":"f877171ef4d77062","repo":"jax-ml/jax","slug":"associative-scan-over-axis-of-non-constant-size","errorCode":null,"errorMessage":"associative scan over axis of non-constant size: {}. You may be able to avoid this on TPU. See b/274176030.","messagePattern":"associative scan over axis of non-constant size: (.+?)\\. You may be able to avoid this on TPU\\. See b/274176030\\.","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2895,"sourceCode":"    raise TypeError(\"lax.associative_scan: fn argument should be callable.\")\n  elems_flat, tree = tree_flatten(elems)\n\n  if reverse:\n    elems_flat = [lax.rev(elem, [axis]) for elem in elems_flat]\n\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`.","sourceCodeStart":2877,"sourceCodeEnd":2913,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2877-L2913","documentation":"associative_scan implements a work-efficient parallel prefix scan whose lowering depends on knowing the scan-axis length at compile (trace) time. If the scanned axis has a non-constant (polynomial/dynamic) size, e.g. under jax.export or with shape-poly dimensions, JAX cannot build the fixed combine tree and raises NotImplementedError, pointing to b/274176030 for TPU workarounds.","triggerScenarios":"Scanning along an axis whose dimension is a symbolic/polynomial dimension (jax.experimental.export, shape polymorphism) or otherwise not statically known, e.g. associative_scan(fn, x) where x.shape[axis] is a DimExpr rather than an int.","commonSituations":"Using jax.export with polymorphic shapes and then calling the exported artifact with dynamic sequence lengths; migrating scan code to JAX shape-polymorphic pipelines; running on TPU where dynamic-size loops are otherwise supported.","solutions":["Make the scanned axis a static dimension (exclude it from polymorphic dims) so its size is known at trace time","Pad/truncate the sequence to a fixed length before the scan","Rewrite as an explicit lax.fori_loop / lax.scan over the dynamic axis, which tolerates non-constant bounds on TPU","Follow b/274176030 for the TPU-specific workaround"],"exampleFix":"// before\nexp.export(fn, polynomial_shapes=('[n,]',))  # n dynamic, then associative_scan over axis 0\n// after\nexp.export(fn, polynomial_shapes=('(_,_)',))  # make scan axis static, pad input to fixed length","handlingStrategy":"validation","validationCode":"import jax\nsize = jax.api shapes... # check statically:\n# x = jnp.ones((n, d))\nassert jnp.asarray(x).shape[axis] is not a dynamic dim — simplest: ensure shape ints known:\nassert isinstance(x.shape[axis], int)","typeGuard":null,"tryCatchPattern":"try:\n    lax.associative_scan(fn, elems)\nexcept NotImplementedError as e:\n    if 'non-constant size' in str(e):\n        out = lax.fori_loop_based_scan(fn, elems)  # fallback explicit loop\n    else:\n        raise","preventionTips":["Keep the scanned axis static when exporting with polymorphic shapes","Pad to fixed lengths for sequence data","Consider lax.scan which supports dynamic lengths"],"tags":["jax","associative-scan","dynamic-shape","shape-polymorphism"],"backgroundTag":"dynamic-shape-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}