{"record":{"id":"c48c15dfbd25c629","repo":"jax-ml/jax","slug":"unsupported-dtype-ref-dtype","errorCode":null,"errorMessage":"Unsupported dtype: {ref.dtype}","messagePattern":"Unsupported dtype: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":495,"sourceCode":"    return x + alignment - rem\n  return x\n\n\n# A tree of `GPUMemoryRef`s.\n_GPUMemoryRefTree = Any\n\n\ndef _ref_group_size(refs: _GPUMemoryRefTree) -> int:\n  size = 0\n  for ref in jax.tree.leaves(refs):\n    # Make sure that the start of each ref is aligned with `SMEM_ALIGNMENT`.\n    size = align_to(size, SMEM_ALIGNMENT)\n    if jnp.issubdtype(ref.dtype, jnp.integer):\n      nbits = jnp.iinfo(ref.dtype).bits\n    elif jnp.issubdtype(ref.dtype, jnp.floating):\n      nbits = jnp.finfo(ref.dtype).bits\n    else:\n      raise NotImplementedError(f\"Unsupported dtype: {ref.dtype}\")\n    ref_bits = math.prod(ref.shape) * nbits\n    if ref_bits % 8:\n      raise ValueError(\n          \"Only byte-aligned shapes are supported. Got shape:\"\n          f\" {ref.dtype}{ref.shape}\"\n      )\n    size += ref_bits // 8\n  return size\n\n\ndef _ref_group_tmem_col_size(refs: _GPUMemoryRefTree) -> int:\n  \"\"\"Returns the total number of TMEM columns used by a group of aliased Refs.\n  \"\"\"\n  ncols = 0\n  for ref in jax.tree.leaves(refs):\n    ref_ncols = ref.layout.cols_in_shape(ref.shape,\n                                         dtypes.itemsize_bits(ref.dtype))\n    ncols += align_to(ref_ncols, TMEM_COL_ALIGNMENT)","sourceCodeStart":477,"sourceCodeEnd":513,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L477-L513","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Change the ref dtype to a supported integer or float dtype (e.g. represent complex64 as two float32 planes)","For bool, store as jnp.int8 or a bitmask and convert in the kernel","File/track an upstream feature request if you need complex support in SMEM allocations"],"exampleFix":"# before\nscratch = pl_core.SMEM((128, 128), jnp.complex64)\n# after (represent complex as two float32 buffers)\nscratch_re = pl_core.SMEM((128, 128), jnp.float32)\nscratch_im = pl_core.SMEM((128, 128), jnp.float32)","handlingStrategy":"type-guard","validationCode":"def check_ref_dtypes(refs):\n    for r in jax.tree.leaves(refs):\n        assert jnp.issubdtype(r.dtype, jnp.integer) or jnp.issubdtype(r.dtype, jnp.floating), r.dtype","typeGuard":"def is_supported_ref_dtype(dt) -> bool:\n    return jnp.issubdtype(dt, jnp.integer) or jnp.issubdtype(dt, jnp.floating)","tryCatchPattern":null,"preventionTips":["Validate all scratch/input dtypes before kernel launch","Represent complex data as paired float buffers"],"tags":["jax","pallas","mosaic-gpu","dtype","shared-memory","not-implemented"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}