jax-ml/jax · error · ValueError

group_offset must be a ()-shaped array. Got: {group_offset.s

Error message

group_offset must be a ()-shaped array. Got: {group_offset.shape}.

What it means

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.

Source

Thrown at jax/experimental/pallas/ops/tpu/megablox/gmm.py:355

    interpret: Whether or not to run the kernel in interpret mode, helpful for
      testing and debugging.

  Returns:
    A 2d, jnp.ndarray with shape [m, n].
  """

  if existing_out is not None:
    assert isinstance(existing_out, jax.Array)
    expected_dtype = existing_out.dtype
    if expected_dtype != preferred_element_type:
      raise ValueError(
          "Existing output dtype must match preferred_element_type."
      )
  if group_offset is None:
    group_offset = jnp.array([0], dtype=jnp.int32)
  else:
    if group_offset.shape:
      raise ValueError(
          f"group_offset must be a ()-shaped array. Got: {group_offset.shape}."
      )
    group_offset = group_offset[None]
  num_current_groups = rhs.shape[0]
  num_total_groups = group_sizes.shape[0]
  lhs, group_sizes, input_dtype = _validate_args(
      lhs=lhs, rhs=rhs, group_sizes=group_sizes
  )

  # Gather shape information.
  m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2])
  if transpose_rhs:
    n = rhs.shape[1]

  # If tiling is callable, look up the problem dimensions in the LUT. If no tuned
  # tile dimensions are available throw an error.
  if callable(tiling):
    tiling = tiling(m, k, n)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
gmm(lhs, rhs, group_sizes, group_offset=jnp.array([2], jnp.int32))
// after
gmm(lhs, rhs, group_sizes, group_offset=jnp.array(2, jnp.int32))

Defensive patterns

Strategy: type-guard

Validate before calling

group_offset = None if group_offset is None else jnp.asarray(group_offset).reshape(())

Type guard

def is_scalar_int(x): return isinstance(x, jax.Array) and x.shape == ()

Prevention

Normalize offsets with reshape(()) at call sites
Never copy internal default literals like jnp.array([0]) into user code

When it happens

Trigger: 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 ().

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f4e84f31d6a6541b. Report an issue: GitHub.