jax-ml/jax · error · ValueError

stack expects at least one operand, got 0.

Error message

stack expects at least one operand, got 0.

What it means

lax.stack (backing jnp.stack) needs at least one array to determine the output shape; an empty operand list raises this ValueError during shape rule evaluation. The result shape is each input's shape with a new axis of length len(operands), which is undefined for zero operands.

Source

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

  dimension_attr = mlir.i64_attr(dimension)
  while len(current_xs) > 1:
    current_xs = [hlo.concatenate(current_xs[i:i+k], dimension_attr)
                  for i in range(0, len(current_xs), k)]
  return current_xs[0]

def _concatenate_lower(ctx, *xs, dimension):
  aval_out, = ctx.avals_out
  out = _concatenate_tree(xs, dimension)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(concatenate_p, _concatenate_lower)

# --- stack and unstack primitives ---

def _stack_shape_rule(*operands, axis):
  if not operands:
    msg = "stack expects at least one operand, got 0."
    raise ValueError(msg)
  if len({op.ndim for op in operands}) != 1:
    msg = "Cannot stack arrays with different numbers of dimensions: got {}."
    raise ValueError(msg.format(", ".join(str(o.shape) for o in operands)))
  if len({op.shape for op in operands}) != 1:
    msg = "All input arrays must have the same shape. Got {}."
    raise ValueError(msg.format(", ".join(str(o.shape) for o in operands)))

  shape = list(operands[0].shape)
  shape.insert(axis, len(operands))
  return tuple(shape)

def _stack_dtype_rule(*operands, axis):
  check_same_dtypes('stack', *operands)
  return operands[0].dtype

def _stack_sharding_rule(*operands, axis):
  non_empty_s = [o.sharding for o in operands if not o.sharding.mesh.empty]
  if not non_empty_s:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard empty input: if not arrays: construct the stacked shape explicitly (insert 0 at the axis) using jnp.zeros
2. Fix upstream generation so at least one element exists (validate dataset/batch size)
3. Handle the empty case at a higher level, e.g. return an empty result of the right dtype/shape

Example fix

# before
out = jnp.stack(parts, axis=0)  # parts == []
# after
out = jnp.stack(parts, axis=0) if parts else jnp.zeros((0,) + item_shape, dtype=dtype)

if not arrays:
    stacked = jnp.zeros((0,) + item_shape, dtype=dtype)
else:
    stacked = jnp.stack(arrays, axis=0)

When it happens

Trigger: jnp.stack([]); lax.stack([], axis=0); stacking a dynamically filtered list that becomes empty.

Common situations: vmap/pmap pipelines where a mapped axis has size 0; collecting per-item results in a loop that never executes; empty-batch edge case in data loaders.

Related errors

• concatenate expects at least one operand, got 0.
• Cannot stack arrays with different numbers of dimensions: go
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must

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