{"record":{"id":"11d4336555dce791","repo":"jax-ml/jax","slug":"ragged-dot-general-requires-dim-name-numbers-to","errorCode":null,"errorMessage":"ragged_dot_general requires {dim_name} numbers to be nonnegative and less than the number of axes of the {arg_name} value, got {dim} for {arg_name} of rank {rank}.","messagePattern":"ragged_dot_general requires (.+?) numbers to be nonnegative and less than the number of axes of the (.+?) value, got (.+?) for (.+?) of rank (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6369,"sourceCode":"      return batch + contract[: contract.index(ragged_dim)]\n    case RaggedDotMode.RAGGED_BATCH:\n      return batch[: batch.index(ragged_dim)]\n\n\ndef _ragged_dot_general_shape_rule(\n    lhs,\n    rhs,\n    group_sizes,\n    *,\n    ragged_dot_dimension_numbers,\n    precision,\n    preferred_element_type: DTypeLike | None,\n    group_offset,\n    out_sharding,\n):\n  def _check_in_range(dim, rank, dim_name, arg_name):\n    if dim < 0 or dim >= rank:\n      raise TypeError(\n          f'ragged_dot_general requires {dim_name} numbers to be nonnegative '\n          f'and less than the number of axes of the {arg_name} value, '\n          f'got {dim} for {arg_name} of rank {rank}.'\n      )\n\n  # Validate the lhs ragged dimension, and find out which mode we're in.\n  if len(ragged_dot_dimension_numbers.lhs_ragged_dimensions) != 1:\n    raise TypeError(\n        'ragged_dot_general expects exactly one lhs ragged dimension.'\n    )\n  lhs_ragged_dim = ragged_dot_dimension_numbers.lhs_ragged_dimensions[0]\n  _check_in_range(lhs_ragged_dim, lhs.ndim, 'lhs ragged dimension', 'lhs')\n  mode = _ragged_dot_mode(lhs.ndim, ragged_dot_dimension_numbers)\n\n  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = (\n      ragged_dot_dimension_numbers.dot_dimension_numbers\n  )\n","sourceCodeStart":6351,"sourceCodeEnd":6387,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6351-L6387","documentation":"Range check inside ragged_dot_general: a supplied dimension index (lhs/rhs ragged dim, or rhs group dim) is negative or >= the rank of the corresponding operand.","triggerScenarios":"jax.lax.ragged_dot_general where lhs_ragged_dim >= lhs.ndim or rhs_group_dim >= rhs.ndim (e.g. using -1 as an index is not allowed, and off-by-one on rank).","commonSituations":"Python-numpy habit of negative indices (not accepted here); shapes changed after adding a batch axis so old dim indices now exceed the rank.","solutions":["Use nonnegative indices strictly less than the operand's .ndim","Recompute all dim indices after any reshape/transpose of lhs/rhs","Add asserts on 0 <= dim < operand.ndim before the call"],"exampleFix":"# before\nout = lax.ragged_dot_general(x, y, gs, dims_with_ragged_dim=-1)\n# after\nout = lax.ragged_dot_general(x, y, gs, dims_with_ragged_dim=x.ndim - 1)","handlingStrategy":"validation","validationCode":"def check_dim(dim, rank):\n    assert 0 <= dim < rank, f'{dim} not in [0, {rank})'\ncheck_dim(lhs_ragged_dim, lhs.ndim)\nfor d in rhs_group_dims: check_dim(d, rhs.ndim)","typeGuard":"def dim_in_range(dim, operand):\n    return isinstance(dim, int) and 0 <= dim < operand.ndim","tryCatchPattern":null,"preventionTips":["No negative indices in ragged dot specs","Recompute indices after adding/removing batch axes"],"tags":["jax","ragged-dot-general","index-out-of-range"],"backgroundTag":"index-out-of-range","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}