jax-ml/jax · error · TypeError

cannot splat a tensor

Error message

cannot splat a tensor

What it means

_splat wraps Triton's splat op, which broadcasts a scalar into a tensor of the given shape; it type-errors if handed a value that is already a RankedTensorType. It is reached from _expand_dims and _reshape when those helpers need to broadcast a value that should be scalar but is actually a tensor.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1543

def _zeros(t: ir.Type) -> ir.Value:
  return _full(t, 0)


def _zeros_like(x: ir.Value) -> ir.Value:
  return _full(x.type, 0)


def _ones(t: ir.Type) -> ir.Value:
  return _full(t, 1)


def _ones_like(x: ir.Value) -> ir.Value:
  return _full(x.type, 1)


def _splat(x: ir.Value, shape: Sequence[int]) -> ir.Value:
  if isinstance(x.type, ir.RankedTensorType):
    raise TypeError("cannot splat a tensor")
  if not shape:
    return x
  return tt_dialect.splat(ir.RankedTensorType.get(shape, x.type), x)


def _expand_dims(x: ir.Value, axis: int) -> ir.Value:
  if not isinstance(x.type, ir.RankedTensorType):
    shape = list(ir.RankedTensorType(x.type).shape)
    shape.insert(axis, 1)
    return _splat(x, shape)
  return tt_dialect.expand_dims(x, axis)


def _float_float_cast(src: ir.Value, dst_type: ir.Type) -> ir.Value:
  src_element_type = ir.FloatType(_element_type(src.type))
  dst_element_type = ir.FloatType(_element_type(dst_type))
  if src_element_type.width == 8 or dst_element_type.width == 8:
    rounding = (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape the kernel logic to keep values tensor-shaped (minimum block dim of 1) rather than scalar-shaped; avoid jnp.reshape(x, ()) inside Triton kernels
2. Upgrade jax (and jax-triton if installed separately) to a version where scalar block handling in expand_dims/reshape lowering is fixed
3. If you control the input, convert the tensor to a scalar before the operation (e.g. extract the element) so _splat sees a scalar

Example fix

// before
s = kernel_ref.reshape(())  # scalar-shaped block hits _splat with tensor
// after
s = kernel_ref.reshape((1,))  # keep 1-element tensor shape

def check_block_shapes(shapes):
    assert all(len(s) > 0 and all(d >= 1 for d in s) for s in shapes), \
        'avoid scalar-shaped (() ) blocks; use (1,) instead'

When it happens

Trigger: Internal lowering paths (_expand_dims/_reshape) receiving a tensor-typed ir.Value where a scalar is expected — typically a consequence of earlier lowering producing a 0-d or 1-d tensor instead of a scalar, e.g. when reshaping a block to a scalar shape or expanding dims of an already-tensorized index value.

Common situations: Pallas kernels that use scalar-shaped blocks (shape ()) or reshape blocks down to scalars; grid/block index arithmetic that yields tensor-typed values; version mismatches between jax and jax-triton where scalar handling changed.

Related errors

• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unfolding dimensions is not supported when commuting an `Un
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Reshape shape {shape} is not divisible by swizzle elements {

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