{"record":{"id":"9c45a894d7e88ce8","repo":"jax-ml/jax","slug":"invalid-axis-axis-for-operand-shape-operand-sha","errorCode":null,"errorMessage":"Invalid axis {axis} for operand shape {operand.shape}","messagePattern":"Invalid axis (.+?) for operand shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8659,"sourceCode":"\ndef _reduce_min_ur_rule(operand, *, axes, out_sharding):\n  out_unreduced, kind = _reduce_op_unreduced_rule(\n      operand, axes, out_sharding, UnreducedKind.min, 'reduce_min')\n  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_min')\n  return out_unreduced, out_reduced, kind\n\nreduce_min_p = standard_primitive(\n    _reduce_op_shape_rule, input_dtype, 'reduce_min',\n    sharding_rule=_reduce_op_sharding_rule_with_out_sharding,\n    vma_rule=partial(core.standard_vma_rule, 'reduce_min'),\n    ur_rule=_reduce_min_ur_rule)\nad.defjvp2(reduce_min_p, _reduce_chooser_jvp_rule)\nbatching.defreducer(reduce_min_p)\n\ndef _argminmax_shape_rule(operand, *, axes, index_dtype):\n  axis, = axes\n  if not (0 <= axis < len(operand.shape)):\n    raise ValueError(f\"Invalid axis {axis} for operand shape {operand.shape}\")\n  if operand.shape[axis] < 1:\n    raise ValueError(\"argmin and argmax require non-empty reduced dimension. \"\n                     f\"operand.shape={operand.shape} {axis=}\")\n  return util.tuple_delete(operand.shape, axis)\n\ndef _argminmax_sharding_rule(operand, *, axes, index_dtype):\n  axis, = axes\n  return operand.sharding.update(spec=\n      util.tuple_delete(operand.sharding.spec, axis))\n\ndef _argminmax_dtype_rule(operand, *, axes, index_dtype):\n  if not dtypes.issubdtype(index_dtype, np.integer):\n    raise TypeError(\"index_dtype must be an integer type, but got {}\"\n                    .format(dtype_to_string(index_dtype)))\n  return index_dtype\n\nclass _ArgMinMaxReducer:\n","sourceCodeStart":8641,"sourceCodeEnd":8677,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8641-L8677","documentation":"A reduction operation (reduce_min/reduce_max family) was given an axis index that is outside the valid range for the operand's rank. The shape rule canonicalizes and validates each entry of the axes argument against operand.shape and rejects out-of-bounds or non-integer axes.","triggerScenarios":"lax.argmin(x, axis=2) on a 2-D array; jnp.argmax(x, axis=x.ndim) after squeeze removed a dimension. The check runs during abstract evaluation, so it fires at trace time.","commonSituations":"Hardcoded axis constants broken by a reshape/squeeze upstream; axis derived from a configuration dict; code ported from NumPy where a later check would have caught it differently.","solutions":["Log operand.shape and axis; fix the axis to a valid one or use -1 for the last axis.","Guard: axis = axis % x.ndim before calling argmin/argmax.","If rank can vary, branch on x.ndim instead of assuming fixed rank."],"exampleFix":"# before\ni = jnp.argmax(x, axis=2)  # x.shape == (8, 4)\n# after\ni = jnp.argmax(x, axis=-1)","handlingStrategy":"validation","validationCode":"axis = axis % x.ndim if isinstance(axis, int) else axis\nassert 0 <= axis < x.ndim, (axis, x.shape)\ni = jnp.argmax(x, axis=axis)","typeGuard":"def valid_axis(x, axis):\n    return -x.ndim <= axis < x.ndim","tryCatchPattern":null,"preventionTips":["Use axis=-1 for the last dimension.","Recheck hardcoded axes after squeeze/vmap changes."],"tags":["jax","argmax","argmin","axis","validation"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}