{"record":{"id":"4bfc778f4202e684","repo":"jax-ml/jax","slug":"please-file-an-issue-at-https-github-com-jax-ml","errorCode":null,"errorMessage":"Please file an issue at https://github.com/jax-ml/jax/issues","messagePattern":"Please file an issue at https://github\\.com/jax-ml/jax/issues","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2647,"sourceCode":"      platform=platform, is_async=is_async)\n\nmlir.register_lowering(all_gather_reduced_p, _all_gather_reduced_lowering)\nfor p in (\"cuda\", \"rocm\", \"tpu\"):\n  mlir.register_lowering(all_gather_reduced_p,\n                         partial(_all_gather_reduced_lowering, platform=p),\n                         platform=p)\n\ndef _all_gather_reduced_transpose_rule(\n    cts, x, *, all_gather_dimension, axis_name, axis_size, tiled):\n  return (unreduced_psum_scatter(cts, axis_name=axis_name,\n                                 scatter_dimension=all_gather_dimension,\n                                 tiled=tiled),)\nad.deflinear2(all_gather_reduced_p, _all_gather_reduced_transpose_rule)\n\ndef _all_gather_reduced_batched_collective(\n    axis_data, vals_in, dims_in, all_gather_dimension, axis_name, axis_size,\n    tiled):\n  raise NotImplementedError(\n      \"Please file an issue at https://github.com/jax-ml/jax/issues\")\nbatching.fancy_primitive_batchers[all_gather_reduced_p] = _all_gather_reduced_batched_collective\n\n####################### unreduced_psum_scatter ###########################\n\n# Unreduced -> Varying collective\ndef unreduced_psum_scatter(x, axis_name, *, scatter_dimension=0, tiled=False):\n  if not isinstance(axis_name, tuple):\n    axis_name = (axis_name,)\n  if not axis_name:\n    return x\n  axis_size = _axis_size(axis_name, None)\n  def bind(leaf):\n    return unreduced_reduce_scatter_p.bind(\n        leaf, axis_name=axis_name, scatter_dimension=scatter_dimension,\n        axis_size=axis_size, tiled=tiled)\n  return tree_util.tree_map(bind, x)\n","sourceCodeStart":2629,"sourceCodeEnd":2665,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2629-L2665","documentation":"`all_gather_reduced` has no batching (vmap) rule registered, so applying `jax.vmap` (or any API that uses fancy batching, like scan with batched carry) over a function containing it raises NotImplementedError with a request to file an issue. The batching rule `_all_gather_reduced_batched_collective` is a stub that always raises. It is a known API gap, not a user-logic error.","triggerScenarios":"Wrapping a function that calls `lax.all_gather_reduced` in `jax.vmap`; batched autodiff or `lax.scan` that triggers the fancy batching machinery on this primitive.","commonSituations":"Ensembling / batched training loops where a per-example function contains a collective; using vmap to auto-batch code written for the named-axis mesh API.","solutions":["Move the collective outside the vmap'd function (manual batching via a leading axis)","Use in_axes=None for the collective arguments or restructure with explicit loop/scan semantics","File an issue at https://github.com/jax-ml/jax/issues as the message requests and pin a JAX version behavior","Check newer JAX versions where a batching rule may have been added"],"exampleFix":"// before\nf = jax.vmap(lambda x: lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8))\ny = f(xs)\n// after\ndef body(x):\n    return lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8)\ny = jax.vmap(body, in_axes=(None,), out_axes=None)(xs)  # or hoist collective out of vmap","handlingStrategy":"fallback","validationCode":"null","typeGuard":"null","tryCatchPattern":"try:\n    ys = jax.vmap(f)(xs)\nexcept NotImplementedError as e:\n    if 'all_gather_reduced' in str(e):\n        ys = jax.lax.map(f, xs)  # sequential fallback\n    else:\n        raise","preventionTips":["Hoist collectives out of vmap'd functions","Check for a registered batching rule before vmap-ing exotic primitives","Watch jax release notes for all_gather_reduced batching support"],"tags":["jax","vmap","collectives","not-implemented","all-gather"],"backgroundTag":"jax-vmap-unsupported-primitive","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}