{"record":{"id":"79b6c8ed1736c65f","repo":"jax-ml/jax","slug":"unsupported-cast-from-dtype-to-dtype","errorCode":null,"errorMessage":"Unsupported cast: {from_dtype} -> {to_dtype}","messagePattern":"Unsupported cast: (.+?) -> (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3149,"sourceCode":"      if from_bitwidth < 32:\n        x = x.astype(jnp.float32)\n    elif jnp.issubdtype(from_dtype, jnp.integer):\n      if from_bitwidth < 32:\n        x = x.astype(jnp.int32)\n    return x != jnp.asarray(0, dtype=x.dtype)\n  if jnp.issubdtype(from_dtype, jnp.signedinteger):\n    if from_bitwidth < 32:\n      x = x.astype(jnp.int32)\n    if jnp.issubdtype(to_dtype, jnp.floating) and to_bitwidth < 32:\n      x = x.astype(jnp.float32)\n    return x.astype(to_dtype)\n  if jnp.issubdtype(from_dtype, jnp.unsignedinteger):\n    if from_bitwidth < 32:\n      x = x.astype(jnp.uint32)\n    # unsigned -> float is unsupported. We fall through and raise at the bottom.\n    if not jnp.issubdtype(to_dtype, jnp.floating):\n      return x.astype(to_dtype)\n  raise NotImplementedError(f\"Unsupported cast: {from_dtype} -> {to_dtype}\")\n\n\n@register_lowering_rule(\n    lax.convert_element_type_p, kernel_types=[*tpu_core.CoreType]\n)\ndef _convert_element_type_lowering_rule(\n    ctx: LoweringRuleContext, x, *, new_dtype, weak_type, sharding\n):\n  del weak_type\n  del sharding\n  out_aval = ctx.avals_out[0]\n  in_aval = ctx.avals_in[0]\n  old_dtype = in_aval.dtype\n  out_type = ctx.aval_to_ir_type(out_aval)\n\n  if old_dtype == new_dtype:\n    return x\n","sourceCodeStart":3131,"sourceCodeEnd":3167,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3131-L3167","documentation":"The dtype-cast helper rejects casts it cannot lower: notably unsigned-to-float casts (uint -> float is unsupported on Mosaic TPU), after already handling unsigned->smaller-unsigned and unsigned widening. Any remaining combination falls through to this NotImplementedError.","triggerScenarios":"convert_element_type from an unsigned integer dtype (e.g. uint8) to a floating dtype (float32/bfloat16) inside a Pallas Mosaic TPU kernel; also any other unhandled from->to pair.","commonSituations":"Preprocessing uint8 image data to float inside a Pallas kernel; normalizing quantized values with .astype(jnp.float32).","solutions":["Cast uint -> int32 first, then int32 -> float: x.astype(jnp.int32).astype(jnp.float32)","Do the float conversion outside the kernel, pass float inputs in","Store source data as signed int8/int32"],"exampleFix":"// before\nxf = x_u8.astype(jnp.float32)\n// after\nxf = x_u8.astype(jnp.int32).astype(jnp.float32)","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\ndef safe_cast(x, to):\n    if jnp.issubdtype(x.dtype, jnp.unsignedinteger) and jnp.issubdtype(to, jnp.floating):\n        return x.astype(jnp.int32).astype(to)\n    return x.astype(to)","typeGuard":"def needs_uint_detour(from_dt, to_dt) -> bool:\n    import jax.numpy as jnp\n    return (jnp.issubdtype(from_dt, jnp.unsignedinteger)\n            and jnp.issubdtype(to_dt, jnp.floating))","tryCatchPattern":null,"preventionTips":["Never cast uint directly to float in Mosaic kernels","Centralize casts in a helper","Convert data to float on the host before passing to kernels"],"tags":["jax","pallas","tpu","dtype-cast","unsigned-integer"],"backgroundTag":"unsupported-dtype-cast","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}