{"record":{"id":"fec444dbaba02ed2","repo":"jax-ml/jax","slug":"argmin-and-argmax-require-non-empty-reduced-dimens","errorCode":null,"errorMessage":"argmin and argmax require non-empty reduced dimension. operand.shape={operand.shape} {axis=}","messagePattern":"argmin and argmax require non-empty reduced dimension\\. operand\\.shape=(.+?) (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8661,"sourceCode":"  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\n  def __init__(self, value_comparator: Callable[[Any, Any], Any]):\n    self._value_comparator = value_comparator","sourceCodeStart":8643,"sourceCodeEnd":8679,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8643-L8679","documentation":"argmin/argmax require the reduced dimension to have size >= 1; reducing an empty dimension has no well-defined index. The check is on the static shape known at trace time.","triggerScenarios":"jnp.argmax(x, axis=0) where x.shape[0] == 0 (e.g., an empty batch after filtering, or a zero-length sequence dimension).","commonSituations":"Data pipelines where a filter/batch can legitimately produce zero rows; sequence models with length-0 sequences; dynamic shapes that statically fold to 0.","solutions":["Check for empty dimensions before calling argmax and handle the empty case separately (return a sentinel or skip).","Fix upstream filtering so the batch is non-empty, or pad with a dummy row before reducing.","If using vmap/jit, add a size guard on the concrete shape via jax.lax.cond or host-side branching."],"exampleFix":"# before\nbest = jnp.argmax(scores, axis=0)  # scores.shape[0] == 0\n# after\nbest = jnp.argmax(scores, axis=0) if scores.shape[0] > 0 else -1","handlingStrategy":"validation","validationCode":"if x.shape[axis] == 0:\n    best = -1  # or skip\nelse:\n    best = jnp.argmax(x, axis=axis)","typeGuard":"def nonempty_axis(x, axis):\n    return x.shape[axis] >= 1","tryCatchPattern":null,"preventionTips":["Guard batch/sequence dims for emptiness before argmax.","Pad or skip empty batches in data pipelines."],"tags":["jax","argmax","empty-array","shape-validation"],"backgroundTag":"reduction-on-empty-array","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}