{"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":"argmin/argmax shape validation found the reduction axis outside the valid range for the operand's shape. The single axis must satisfy 0 <= axis < len(operand.shape).","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","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}