{"record":{"id":"aa7a830af7982a63","repo":"jax-ml/jax","slug":"ragged-dot-general-expects-exactly-one-lhs-ragged","errorCode":null,"errorMessage":"ragged_dot_general expects exactly one lhs ragged dimension.","messagePattern":"ragged_dot_general expects exactly one lhs ragged dimension\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6377,"sourceCode":"    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\n  # Validate the shape of group_sizes, if it is something other than [g].\n  if group_sizes.ndim == 0:\n    raise TypeError('expected rank of group_sizes to be >=1.')\n  if group_sizes.ndim != 1:\n    # Construct the expected shape [b...,x...,g] of group_sizes.\n    prefix_dims = _ragged_dot_prefix_dims(\n        mode, lhs.ndim, lhs_ragged_dim, lhs_batch, lhs_contracting\n    )","sourceCodeStart":6359,"sourceCodeEnd":6395,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6359-L6395","documentation":"ragged_dot_general only supports exactly one ragged dimension on the lhs; this fires when len(lhs_ragged_dimensions) != 1 (zero or multiple entries).","triggerScenarios":"Passing RaggedDotDimensionNumbers with lhs_ragged_dimensions=[] or with two indices, e.g. (1, 2), to jax.lax.ragged_dot_general.","commonSituations":"Migrating from XLA RaggedDot which may allow multiple ragged dims; misunderstanding that only the lhs carries one ragged dim (rhs raggedness comes via group dims).","solutions":["Supply exactly one lhs ragged dimension; model multiple ragged axes with separate ragged_dot_general calls or reshapes"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(rdn.lhs_ragged_dimensions) == 1, 'exactly one lhs ragged dim required'","typeGuard":"def single_ragged_dim(rdn):\n    return len(rdn.lhs_ragged_dimensions) == 1","tryCatchPattern":null,"preventionTips":["Model only the lhs as ragged; chain calls for multiple ragged axes"],"tags":["jax","ragged-dot-general","dimension-numbers"],"backgroundTag":"invalid-dimension-numbers","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}