jax-ml/jax · error · NotImplementedError

Unsupported dtype: {ref.dtype}

Error message

Unsupported dtype: {ref.dtype}

What it means

When Mosaic GPU computes the byte size of a shared-memory (or TMEM) reference group in `_ref_group_size`, it only knows how to count bits for integer and floating dtypes via `jnp.iinfo`/`jnp.finfo`. Any other dtype (complex, bool, extended/bfloat edge cases not covered, custom types) raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:495

    return x + alignment - rem
  return x


# A tree of `GPUMemoryRef`s.
_GPUMemoryRefTree = Any


def _ref_group_size(refs: _GPUMemoryRefTree) -> int:
  size = 0
  for ref in jax.tree.leaves(refs):
    # Make sure that the start of each ref is aligned with `SMEM_ALIGNMENT`.
    size = align_to(size, SMEM_ALIGNMENT)
    if jnp.issubdtype(ref.dtype, jnp.integer):
      nbits = jnp.iinfo(ref.dtype).bits
    elif jnp.issubdtype(ref.dtype, jnp.floating):
      nbits = jnp.finfo(ref.dtype).bits
    else:
      raise NotImplementedError(f"Unsupported dtype: {ref.dtype}")
    ref_bits = math.prod(ref.shape) * nbits
    if ref_bits % 8:
      raise ValueError(
          "Only byte-aligned shapes are supported. Got shape:"
          f" {ref.dtype}{ref.shape}"
      )
    size += ref_bits // 8
  return size


def _ref_group_tmem_col_size(refs: _GPUMemoryRefTree) -> int:
  """Returns the total number of TMEM columns used by a group of aliased Refs.
  """
  ncols = 0
  for ref in jax.tree.leaves(refs):
    ref_ncols = ref.layout.cols_in_shape(ref.shape,
                                         dtypes.itemsize_bits(ref.dtype))
    ncols += align_to(ref_ncols, TMEM_COL_ALIGNMENT)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the ref dtype to a supported integer or float dtype (e.g. represent complex64 as two float32 planes)
2. For bool, store as jnp.int8 or a bitmask and convert in the kernel
3. File/track an upstream feature request if you need complex support in SMEM allocations

Example fix

# before
scratch = pl_core.SMEM((128, 128), jnp.complex64)
# after (represent complex as two float32 buffers)
scratch_re = pl_core.SMEM((128, 128), jnp.float32)
scratch_im = pl_core.SMEM((128, 128), jnp.float32)

def check_ref_dtypes(refs):
    for r in jax.tree.leaves(refs):
        assert jnp.issubdtype(r.dtype, jnp.integer) or jnp.issubdtype(r.dtype, jnp.floating), r.dtype

def is_supported_ref_dtype(dt) -> bool:
    return jnp.issubdtype(dt, jnp.integer) or jnp.issubdtype(dt, jnp.floating)

When it happens

Trigger: Declaring a Pallas kernel scratch or input ref whose dtype is complex (e.g. `jnp.complex64`), bool, or a non-numeric/custom dtype, causing _ref_group_size to fall through both issubtype checks.

Common situations: Porting FFT-style or signal-processing kernels that use complex64 to Mosaic GPU; using bool predicate buffers in SMEM; assuming all jnp dtypes are supported because they work elsewhere in JAX.

Related errors

• Sparse metadata format not implemented for {operand_dtype=}
• Data type bitcast is only supported from i8 to other types.
• out_dtype argument in binary_op_lowering_rule_wg
• {ctx.prim} does not support {x_aval.dtype} and {y_aval.dtype
• bitcast_convert_type with different bitwidths not supported

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