{"record":{"id":"23866ea8f38cc3de","repo":"jax-ml/jax","slug":"expected-rank-of-group-sizes-to-be-1","errorCode":null,"errorMessage":"expected rank of group_sizes to be >=1.","messagePattern":"expected rank of group_sizes to be >=1\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6390,"sourceCode":"          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    )\n    expected_gs_shape = tuple(lhs.shape[i] for i in prefix_dims)\n    expected_gs_shape += (group_sizes.shape[-1],)\n    # TODO(pravnar): Permit other broadcastable shapes.\n    if not core.definitely_equal_shape(group_sizes.shape, expected_gs_shape):\n      raise TypeError(\n          'expected group_sizes to have shape '\n          f'{expected_gs_shape}, got {group_sizes.shape}.'\n      )\n  num_groups = group_sizes.shape[-1]\n  if (mode in (RaggedDotMode.RAGGED_CONTRACTING,\n               RaggedDotMode.RAGGED_NONCONTRACTING)\n      and core.is_symbolic_dim(num_groups)):\n    raise TypeError(","sourceCodeStart":6372,"sourceCodeEnd":6408,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6372-L6408","documentation":"Shape validation for ragged_dot_general's group_sizes: a 0-D (scalar) group_sizes was passed, but rank >= 1 is required. group_sizes must at least be a 1-D vector of per-group sizes [g].","triggerScenarios":"jax.lax.ragged_dot_general(..., group_sizes=jnp.array(5)) or a Python int/0-D tracer for group_sizes.","commonSituations":"Single-group use where the developer passes a scalar total instead of [total]; forgetting jnp.array wrapping around a list.","solutions":["Wrap the scalar into a 1-D array: group_sizes=jnp.array([n])","For batched/nested cases, give group_sizes shape [b..., x..., g] matching the ragged-dot prefix dims"],"exampleFix":"# before\nout = lax.ragged_dot_general(x, y, jnp.array(10), dims)\n# after\nout = lax.ragged_dot_general(x, y, jnp.array([10]), dims)","handlingStrategy":"validation","validationCode":"assert group_sizes.ndim >= 1, 'group_sizes must be at least 1-D'\nif np.ndim(group_sizes) == 0:\n    group_sizes = jnp.array([group_sizes])","typeGuard":"def valid_group_sizes(gs):\n    return getattr(gs, 'ndim', np.ndim(gs)) >= 1","tryCatchPattern":null,"preventionTips":["Always construct group_sizes with jnp.array([...])","Pass [total] for the single-group case"],"tags":["jax","ragged-dot-general","group-sizes","shape-mismatch"],"backgroundTag":"invalid-argument-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}