jax-ml/jax · error · NotImplementedError

Unimplemented group_offset support.

Error message

Unimplemented group_offset support.

What it means

The JVP (forward-mode autodiff) rule for ragged_dot_general does not support the group_offset argument; it only runs when group_offset is None. Passing a non-None group_offset under jvp/forward-mode raises NotImplementedError.

Source

Thrown at jax/_src/lax/lax.py:6490

  return _dot_general_dtype_rule(
      lhs,
      rhs,
      dimension_numbers=ragged_dot_dimension_numbers.dot_dimension_numbers,
      precision=precision,
      preferred_element_type=preferred_element_type,
      out_sharding=None,
      name='lax.ragged_dot_general',
  )


def _ragged_dot_general_jvp_rule(
    primals, tangents, ragged_dot_dimension_numbers,
    precision, preferred_element_type, group_offset, out_sharding
):
  # note - we could ostensibly just get this by passing on the
  # value to ragged_dot below, but, this feels cleaner.
  if group_offset is not None:
    raise NotImplementedError('Unimplemented group_offset support.')
  x, y, gs = primals
  dx, dy, _ = tangents  # no tan on the gs

  # primal
  primal_out = ragged_dot_general(
      x,
      y,
      gs,
      ragged_dot_dimension_numbers=ragged_dot_dimension_numbers,
      precision=precision,
      preferred_element_type=preferred_element_type,
  )

  # tangent
  dx_out = (
      ragged_dot_general(
          dx,
          y,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the group_offset argument (pass None) when computing forward-mode derivatives
2. Use reverse-mode (jax.grad) instead, if the transpose rule supports group_offset in your JAX version
3. Compute the jvp manually by splitting the computation at offset boundaries

Example fix

// before
out = jax.jvp(lambda x: ragged_dot_general(x, w, gs, dn, mode=mode, group_offset=off), (x,), (dx,))
// after
out = jax.jvp(lambda x: ragged_dot_general(x, w, gs, dn, mode=mode), (x,), (dx,))

if group_offset is not None and using_jvp:
    group_offset = None  # restructure instead

except NotImplementedError as e:
    if 'group_offset' in str(e): group_offset = None; recompute()

When it happens

Trigger: Using jax.jvp (directly or via libraries that use forward-mode, e.g. jax.checkpoint in some modes or IMC/forward-over-reverse setups) on a function that calls ragged_dot_general with group_offset set.

Common situations: Using group_offset to resume from a previous batch segment (multi-step MoE dispatch) and then differentiating with jvp; upgrading code that worked without offsets.

Related errors

• for grad support, subclass {type(self)} must implement `vjp_
• for jvp support, subclass {type(self)} must implement `jvp`
• for linearize support, subclass {type(self)} must implement
• for transpose support, subclass {type(self)} must implement
• open an issue at https://github.com/google/jax !!

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