{"record":{"id":"0715a68b72d71c25","repo":"jax-ml/jax","slug":"only-arrays-with-32-bit-element-types-can-be-conve","errorCode":null,"errorMessage":"Only arrays with 32-bit element types can be converted to scalars, but got: {aval_out.dtype}. Try casting the input before squeezing the scalar.","messagePattern":"Only arrays with 32-bit element types can be converted to scalars, but got: (.+?)\\. Try casting the input before squeezing the scalar\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3237,"sourceCode":"  if dimensions is not None:\n    raise NotImplementedError\n  if any(d is None for d in new_sizes):\n    raise NotImplementedError\n  if not ctx.avals_in[0].shape:\n    return vector.broadcast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)\n  if not ctx.avals_out[0].shape:\n    return vector.extract(x, [], [0] * len(ctx.avals_in[0].shape))\n  return vector.shape_cast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)\n\n\n@register_lowering_rule(lax.squeeze_p, kernel_types=[*tpu_core.CoreType])\ndef _squeeze_lowering_rule(ctx: LoweringRuleContext, x, dimensions):\n  del dimensions  # Unused.\n  (aval_in,) = ctx.avals_in\n  (aval_out,) = ctx.avals_out\n  if not aval_out.shape:\n    if aval_out.dtype.itemsize != 4:\n      raise ValueError(\n          \"Only arrays with 32-bit element types can be converted to scalars,\"\n          f\" but got: {aval_out.dtype}. Try casting the input before squeezing\"\n          \" the scalar.\"\n      )\n    return vector.extract(x, [], [0] * len(aval_in.shape))\n  return vector.shape_cast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)\n\n\n@register_lowering_rule(lax.concatenate_p, kernel_types=[*tpu_core.CoreType])\ndef _concatenate_lowering_rule(ctx: LoweringRuleContext, *xs, dimension):\n  del ctx  # Unused.\n  return tpu.concatenate(xs, dimension=dimension)\n\n@register_lowering_rule(lax_internal.stack_p, kernel_types=[*tpu_core.CoreType])\ndef _stack_lowering_rule(ctx: LoweringRuleContext, *xs, axis):\n  x_aval = ctx.avals_in[0]\n\n  new_shape = list(x_aval.shape)","sourceCodeStart":3219,"sourceCodeEnd":3255,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3219-L3255","documentation":"Squeezing an array down to a scalar (empty output shape) in Mosaic requires a 32-bit element type, because the lowering extracts via vector ops that only handle 4-byte scalars. Wider/narrower dtypes raise ValueError with a suggestion to cast first.","triggerScenarios":"lax.squeeze (or jnp.squeeze) producing a 0-d scalar whose dtype itemsize != 4 (e.g. float64, int16, bfloat16? — any non-4-byte dtype) inside a Pallas Mosaic TPU kernel.","commonSituations":"Reducing a block to a scalar (e.g. loss value) in a kernel using float64 or sub-32-bit dtypes; x64-enabled JAX making intermediates 64-bit.","solutions":["Cast to a 32-bit dtype before squeezing: x.astype(jnp.float32) then squeeze","Avoid producing scalars in-kernel; keep a shape-(1,) result","Disable jax_enable_x64 so values default to float32/int32"],"exampleFix":"// before\ns = jnp.squeeze(loss_f64)\n// after\ns = jnp.squeeze(loss_f64.astype(jnp.float32))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nif not out.shape:\n    assert out.dtype.itemsize == 4, 'cast to 32-bit before squeezing scalar'","typeGuard":"def squeeze_ok(aval) -> bool:\n    return bool(aval.shape) or aval.dtype.itemsize == 4","tryCatchPattern":null,"preventionTips":["Cast scalars to float32/int32 before squeeze","Keep reduced values as shape-(1,) arrays in kernels","Disable x64 in TPU sessions"],"tags":["jax","pallas","tpu","squeeze","scalar","dtype"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}