jax-ml/jax · error · ValueError

Only arrays with 32-bit element types can be converted to sc

Error message

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.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3237

  if dimensions is not None:
    raise NotImplementedError
  if any(d is None for d in new_sizes):
    raise NotImplementedError
  if not ctx.avals_in[0].shape:
    return vector.broadcast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)
  if not ctx.avals_out[0].shape:
    return vector.extract(x, [], [0] * len(ctx.avals_in[0].shape))
  return vector.shape_cast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)


@register_lowering_rule(lax.squeeze_p, kernel_types=[*tpu_core.CoreType])
def _squeeze_lowering_rule(ctx: LoweringRuleContext, x, dimensions):
  del dimensions  # Unused.
  (aval_in,) = ctx.avals_in
  (aval_out,) = ctx.avals_out
  if not aval_out.shape:
    if aval_out.dtype.itemsize != 4:
      raise ValueError(
          "Only arrays with 32-bit element types can be converted to scalars,"
          f" but got: {aval_out.dtype}. Try casting the input before squeezing"
          " the scalar."
      )
    return vector.extract(x, [], [0] * len(aval_in.shape))
  return vector.shape_cast(ctx.aval_to_ir_type(ctx.avals_out[0]), x)


@register_lowering_rule(lax.concatenate_p, kernel_types=[*tpu_core.CoreType])
def _concatenate_lowering_rule(ctx: LoweringRuleContext, *xs, dimension):
  del ctx  # Unused.
  return tpu.concatenate(xs, dimension=dimension)

@register_lowering_rule(lax_internal.stack_p, kernel_types=[*tpu_core.CoreType])
def _stack_lowering_rule(ctx: LoweringRuleContext, *xs, axis):
  x_aval = ctx.avals_in[0]

  new_shape = list(x_aval.shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to a 32-bit dtype before squeezing: x.astype(jnp.float32) then squeeze
2. Avoid producing scalars in-kernel; keep a shape-(1,) result
3. Disable jax_enable_x64 so values default to float32/int32

Example fix

// before
s = jnp.squeeze(loss_f64)
// after
s = jnp.squeeze(loss_f64.astype(jnp.float32))

import jax.numpy as jnp
if not out.shape:
    assert out.dtype.itemsize == 4, 'cast to 32-bit before squeezing scalar'

def squeeze_ok(aval) -> bool:
    return bool(aval.shape) or aval.dtype.itemsize == 4

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Cannot swap scalars to VMEM.
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.

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