{"record":{"id":"955afc7f615571ff","repo":"jax-ml/jax","slug":"named-reductions-not-implemented-for-jnp-name","errorCode":null,"errorMessage":"Named reductions not implemented for jnp.{name}()","messagePattern":"Named reductions not implemented for jnp\\.(.+?)\\(\\)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/reductions.py","lineNumber":153,"sourceCode":"    if promote_integers:\n      result_dtype = _promote_integer_dtype(result_dtype)\n  else:\n    result_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, name)\n\n  if upcast_f16_for_computation and dtypes.issubdtype(result_dtype, np.inexact):\n    computation_dtype = _upcast_f16(result_dtype)\n  else:\n    computation_dtype = result_dtype\n  a = lax.convert_element_type(a, computation_dtype)\n  op = op if computation_dtype != np.bool_ else bool_op\n  # NB: in XLA, init_val must be an identity for the op, so the user-specified\n  # initial value must be applied afterward.\n  init_val = _reduction_init_val(a, init_val)\n  if where_ is not None:\n    a = _where(where_, a, init_val)\n  if pos_dims is not dims:\n    if parallel_reduce is None:\n      raise NotImplementedError(f\"Named reductions not implemented for jnp.{name}()\")\n    result = parallel_reduce(a, dims)\n  else:\n    result = lax.reduce(a, init_val, op, dims)\n  if initial is not None:\n    initial_arr = lax.convert_element_type(initial, lax.asarray(a).dtype)\n    if initial_arr.shape != ():\n      raise ValueError(\"initial value must be a scalar. \"\n                       f\"Got array of shape {initial_arr.shape}\")\n    result = op(initial_arr, result)\n  if keepdims:\n    result = lax.expand_dims(result, pos_dims)\n  return lax.convert_element_type(result, dtype or result_dtype)\n\ndef _canonicalize_axis_allow_named(x, rank):\n  return maybe_named_axis(x, lambda i: canonicalize_axis(i, rank), lambda name: name)\n\ndef _reduction_dims(a: ArrayLike, axis: Axis):\n  if axis is None:","sourceCodeStart":135,"sourceCodeEnd":171,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/reductions.py#L135-L171","documentation":"The reduction was invoked with named axes (e.g. axis='batch' from pmap/jit named axes) but this particular reduction has no named-axis (collective/pmap) implementation registered in JAX.","triggerScenarios":"Calling jnp.<reduce>(x, axis='dev') inside jax.pmap where the reduction lacks a parallel_reduce implementation; mixing named axes with an uncommon reduction op.","commonSituations":"Using pmap with named axes and reductions that only support positional axes; migrating multi-device code that worked with lax.psum but not this jnp reduction.","solutions":["Use positional integer axes instead of named axes","Use lax.psum / the collective equivalent (e.g. jax.lax.pmax) for named-axis reductions","Move the reduction off the named axis: reduce positional axes first, then use psum on the remainder"],"exampleFix":"// before\njax.pmap(lambda x: jnp.logsumexp(x, axis='dev'), axis_name='dev')(x)\n// after\njax.pmap(lambda x: jax.lax.pmax(x, 'dev').sum(-1), axis_name='dev')(x)","handlingStrategy":"fallback","validationCode":"if isinstance(axis, str):  # named axis unsupported\n    axis = None  # or use lax.psum/pmax collectives instead","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use lax.psum/pmax for named-axis reductions","Prefer integer axes in pmap code"],"tags":["jax","named-axes","pmap","reductions"],"backgroundTag":"named-axis-reduction-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}