{"record":{"id":"f4e84f31d6a6541b","repo":"jax-ml/jax","slug":"group-offset-must-be-a-shaped-array-got-grou","errorCode":null,"errorMessage":"group_offset must be a ()-shaped array. Got: {group_offset.shape}.","messagePattern":"group_offset must be a \\(\\)-shaped array\\. Got: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/megablox/gmm.py","lineNumber":355,"sourceCode":"    interpret: Whether or not to run the kernel in interpret mode, helpful for\n      testing and debugging.\n\n  Returns:\n    A 2d, jnp.ndarray with shape [m, n].\n  \"\"\"\n\n  if existing_out is not None:\n    assert isinstance(existing_out, jax.Array)\n    expected_dtype = existing_out.dtype\n    if expected_dtype != preferred_element_type:\n      raise ValueError(\n          \"Existing output dtype must match preferred_element_type.\"\n      )\n  if group_offset is None:\n    group_offset = jnp.array([0], dtype=jnp.int32)\n  else:\n    if group_offset.shape:\n      raise ValueError(\n          f\"group_offset must be a ()-shaped array. Got: {group_offset.shape}.\"\n      )\n    group_offset = group_offset[None]\n  num_current_groups = rhs.shape[0]\n  num_total_groups = group_sizes.shape[0]\n  lhs, group_sizes, input_dtype = _validate_args(\n      lhs=lhs, rhs=rhs, group_sizes=group_sizes\n  )\n\n  # Gather shape information.\n  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])\n  if transpose_rhs:\n    n = rhs.shape[1]\n\n  # If tiling is callable, look up the problem dimensions in the LUT. If no tuned\n  # tile dimensions are available throw an error.\n  if callable(tiling):\n    tiling = tiling(m, k, n)","sourceCodeStart":337,"sourceCodeEnd":373,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/megablox/gmm.py#L337-L373","documentation":"The group_offset argument to megablox gmm must be a scalar (()-shaped) array; passing a 1-element vector like jnp.array([0]) is rejected. The kernel indexes a single group offset, so any non-empty shape is invalid. The code itself wraps the scalar with [None] after validation, confirming a 0-d input is expected.","triggerScenarios":"Calling gmm with group_offset=jnp.array([2]) (shape (1,)) or any array with ndim > 0. Correct usage is group_offset=jnp.array(2) with shape ().","commonSituations":"Copying the internal default pattern jnp.array([0], dtype=jnp.int32) seen in the same function's source; migrating from an older API that accepted shape-(1,) offsets.","solutions":["Pass a true scalar: group_offset=jnp.array(2, dtype=jnp.int32) (shape ())","Pass group_offset=None to use the default offset of 0"],"exampleFix":"// before\ngmm(lhs, rhs, group_sizes, group_offset=jnp.array([2], jnp.int32))\n// after\ngmm(lhs, rhs, group_sizes, group_offset=jnp.array(2, jnp.int32))","handlingStrategy":"type-guard","validationCode":"group_offset = None if group_offset is None else jnp.asarray(group_offset).reshape(())","typeGuard":"def is_scalar_int(x): return isinstance(x, jax.Array) and x.shape == ()","tryCatchPattern":null,"preventionTips":["Normalize offsets with reshape(()) at call sites","Never copy internal default literals like jnp.array([0]) into user code"],"tags":["jax","pallas","tpu","megablox","shape-validation"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}