jax-ml/jax · error · ValueError

unstack requires arrays with rank > 0, however a scalar arra

Error message

unstack requires arrays with rank > 0, however a scalar array of shape {} was passed.

What it means

lax.unstack splits an array into `shape[axis]` results along the given axis; a 0-D scalar has no axis to split, so the shape rule raises this ValueError. unstack is essentially the inverse of stack and requires rank > 0.

Source

Thrown at jax/_src/lax/lax.py:7439

def _stack_ur_rule(*operands, **kwargs):
  out_unreduced, kind = _concatenate_unreduced_rule(*operands, **kwargs)
  out_reduced = _concatenate_reduced_rule(*operands, **kwargs)
  return out_unreduced, out_reduced, kind

stack_p = standard_primitive(
    _stack_shape_rule, _stack_dtype_rule, 'stack',
    sharding_rule=_stack_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'stack'),
    ur_rule=_stack_ur_rule)
ad.deflinear2(stack_p, _stack_transpose_rule)
mlir.register_lowering(stack_p, _stack_lower)


def _unstack_shape_rule(operand, *, axis):
  if operand.ndim == 0:
    msg = "unstack requires arrays with rank > 0, however a scalar array of shape {} was passed."
    raise ValueError(msg.format(operand.shape))
  shape = list(operand.shape)
  num_results = shape.pop(axis)
  return (tuple(shape),) * num_results

def _unstack_dtype_rule(operand, *, axis):
  num_results = operand.shape[axis]
  return (operand.dtype,) * num_results

def _unstack_weak_type_rule(operand, *, axis):
  num_results = operand.shape[axis]
  return (operand.weak_type,) * num_results

def _unstack_sharding_rule(operand, *, axis):
  if operand.sharding.spec[axis] is not None:
    raise core.ShardingTypeError(
        f"unstack operand cannot be sharded on the unstacking axis {axis}. "
        f"Got operand type={operand.str_short(True)}"
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the operand has rank >= 1: keep dims with keepdims=True in reductions, or add an axis via x[None]
2. Check x.ndim before calling unstack and skip/handle scalars separately
3. Avoid over-squeezing upstream: prefer jnp.squeeze(x, axis=specific_axis)

Example fix

# before
outs = jax.lax.unstack(scalar_loss)  # shape ()
# after
outs = jax.lax.unstack(losses)  # shape (n,) from .sum(axis=...) without squeeze

assert x.ndim > 0, f'unstack needs rank>0, got {x.shape}'

def unstackable(x) -> bool:
    return getattr(x, 'ndim', 0) > 0

When it happens

Trigger: jax.lax.unstack(jnp.float32(3.0)) — unstacking a scalar; passing a value that was over-squeezed; axis computed on data whose rank collapsed to 0.

Common situations: Unstacking a loss or metric that is scalar after mean(); aggressive squeeze() removing all axes; feeding scalars into code that expects batched tensors.

Related errors

• top_k operand must have >= 1 dimension, got {}
• iteration over a 0-d array
• Invalid scalar value {x}
• cond_fun must return a boolean scalar, but got output type(s
• length of padding_config must equal the number of axes of op

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