jax-ml/jax · error · NotImplementedError

Only 1D vectors are supported {v_ty}

Error message

Only 1D vectors are supported {v_ty}

What it means

vector_slice uses extract_strided_slice which only applies along a single dimension, so only rank-1 vectors are supported. Passing a 2D+ vector raises NotImplementedError with the offending vector type in the message.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2110

      raise ValueError(f"Can't bitcast {x.type} to {new_type}")
    return vector.bitcast(new_type, x)
  if isinstance(x.type, ir.IntegerType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.IntegerType):
    return arith.bitcast(new_type, x)
  if isinstance(x.type, ir.FloatType) and isinstance(new_type, ir.FloatType):
    return arith.bitcast(new_type, x)
  raise ValueError(f"Can't bitcast {x.type} to {new_type}")


def ceil_div(x: int, y: int):
  return (x + y - 1) // y


def vector_slice(v: ir.Value, s: slice):
  v_ty = ir.VectorType(v.type)
  if len(v_ty.shape) != 1:
    raise NotImplementedError(f"Only 1D vectors are supported {v_ty}")
  [v_len] = v_ty.shape
  slice_length = len(range(v_len)[s])
  return vector.extract_strided_slice(
      ir.VectorType.get((slice_length,), v_ty.element_type),
      v,
      [s.start or 0],
      [slice_length],
      [1],
  )


def vector_concat(
    vectors: Sequence[ir.Value[ir.VectorType]],
) -> ir.Value[ir.VectorType]:
  if not vectors:
    raise ValueError("Cannot concatenate an empty list of vectors")
  vty = vectors[0].type
  if not isinstance(vty, ir.VectorType):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten to 1D before slicing (vector.shape_cast to rank-1, slice, then cast back)
2. Use extract_strided_slice directly with per-dimension offsets for >=2D vectors
3. Restructure the kernel to keep vectors rank-1 where slicing is needed

Example fix

# before
part = vector_slice(v_2d, slice(0, 2))
# after
flat = vector.shape_cast(ir.VectorType.get((32,), f32), v_2d)
part = vector_slice(flat, slice(0, 16))

assert ir.VectorType(v.type).rank == 1, f'vector_slice needs 1D, got {v.type}'

def is_1d_vector(v) -> bool:
    t = v.type
    return isinstance(t, ir.VectorType) and ir.VectorType(t).rank == 1

When it happens

Trigger: Calling vector_slice(v, s) where v is e.g. of type vector<4x8xf32> — any rank >= 2 vector value.

Common situations: Slicing a matrix tile produced by a 2D tiling layout; upgrading layouts from 1D to 2D without updating slicing logic.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Expected metadata to be 3-dimensional (M, K // 4, 2), but it

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