jax-ml/jax · error · NotImplementedError

ragged_dot vmap over any dim but 0 - NYI

Error message

ragged_dot vmap over any dim but 0 - NYI

What it means

The vmap batch rule for ragged_dot only supports batching over dimension 0 of every batched operand. Batching over any other dimension is not yet implemented (NYI).

Source

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

      if ad.is_undefined_primal(y)
      else _ragged_dot_grad(ct, y, grad_x_dims, x.aval)
  )
  y_bar = (
      None
      if ad.is_undefined_primal(x)
      else _ragged_dot_grad(x, ct, grad_y_dims, y.aval)
  )
  return x_bar, y_bar, None


def _ragged_dot_batch_unpack_args(batched_args):
  lhs, rhs, _ = batched_args
  return (lhs, rhs)


def _ragged_dot_batch_unpack_dims(batch_dims):
  if not all(dim == 0 for dim in batch_dims):
    raise NotImplementedError('ragged_dot vmap over any dim but 0 - NYI')
  lbd, rbd, _ = batch_dims
  return (lbd, rbd)


def _ragged_dot_general_invoke_prim(
    group_sizes,
    lhs,
    rhs,
    new_ragged_dot_dimension_numbers,
    precision,
    preferred_element_type,
    out_sharding,
):
  del out_sharding
  return ragged_dot_general(
      lhs,
      rhs,
      group_sizes,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose/squeeze so the batched dimension is axis 0 before vmap (e.g. x.transpose(1,0,2) or in_axes=0)
2. Use in_axes=0 (or None) for all ragged_dot operands
3. Fall back to a manual loop or lax.map over the batch

Example fix

// before
f = jax.vmap(ragged_step, in_axes=(1, None, None))  # x batched on dim 1
// after
x2 = x.transpose(1, 0, 2)
f = jax.vmap(ragged_step, in_axes=(0, None, None))

assert all(ax in (0, None) for ax in in_axes), 'batch on dim 0 only'

When it happens

Trigger: Applying jax.vmap to a function using ragged_dot/ragged_dot_general where the batched axis maps to a non-zero dimension of lhs, rhs, or group_sizes.

Common situations: Vectorizing a per-example ragged matmul but in_axes points at dim 1 (e.g. in_axes=1 on a (seq, batch, d) tensor), or transposing inputs so the batch dim isn't leading.

Related errors

• reduce_window batching is not implemented for initial values
• Batching over custom allocations is not supported yet.
• Batching over dynamic grid values is not supported yet.
• vmapping pallas_call with no arguments.
• No support for vmapping over nontrivial slices just yet: {id

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