{"record":{"id":"4dbba72e2ba2b3e6","repo":"jax-ml/jax","slug":"unstack-requires-arrays-with-rank-0-however-a-s","errorCode":null,"errorMessage":"unstack requires arrays with rank > 0, however a scalar array of shape {} was passed.","messagePattern":"unstack requires arrays with rank > 0, however a scalar array of shape (.+?) was passed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7439,"sourceCode":"\ndef _stack_ur_rule(*operands, **kwargs):\n  out_unreduced, kind = _concatenate_unreduced_rule(*operands, **kwargs)\n  out_reduced = _concatenate_reduced_rule(*operands, **kwargs)\n  return out_unreduced, out_reduced, kind\n\nstack_p = standard_primitive(\n    _stack_shape_rule, _stack_dtype_rule, 'stack',\n    sharding_rule=_stack_sharding_rule,\n    vma_rule=partial(core.standard_vma_rule, 'stack'),\n    ur_rule=_stack_ur_rule)\nad.deflinear2(stack_p, _stack_transpose_rule)\nmlir.register_lowering(stack_p, _stack_lower)\n\n\ndef _unstack_shape_rule(operand, *, axis):\n  if operand.ndim == 0:\n    msg = \"unstack requires arrays with rank > 0, however a scalar array of shape {} was passed.\"\n    raise ValueError(msg.format(operand.shape))\n  shape = list(operand.shape)\n  num_results = shape.pop(axis)\n  return (tuple(shape),) * num_results\n\ndef _unstack_dtype_rule(operand, *, axis):\n  num_results = operand.shape[axis]\n  return (operand.dtype,) * num_results\n\ndef _unstack_weak_type_rule(operand, *, axis):\n  num_results = operand.shape[axis]\n  return (operand.weak_type,) * num_results\n\ndef _unstack_sharding_rule(operand, *, axis):\n  if operand.sharding.spec[axis] is not None:\n    raise core.ShardingTypeError(\n        f\"unstack operand cannot be sharded on the unstacking axis {axis}. \"\n        f\"Got operand type={operand.str_short(True)}\"\n    )","sourceCodeStart":7421,"sourceCodeEnd":7457,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7421-L7457","documentation":"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.","triggerScenarios":"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.","commonSituations":"Unstacking a loss or metric that is scalar after mean(); aggressive squeeze() removing all axes; feeding scalars into code that expects batched tensors.","solutions":["Ensure the operand has rank >= 1: keep dims with keepdims=True in reductions, or add an axis via x[None]","Check x.ndim before calling unstack and skip/handle scalars separately","Avoid over-squeezing upstream: prefer jnp.squeeze(x, axis=specific_axis)"],"exampleFix":"# before\nouts = jax.lax.unstack(scalar_loss)  # shape ()\n# after\nouts = jax.lax.unstack(losses)  # shape (n,) from .sum(axis=...) without squeeze","handlingStrategy":"type-guard","validationCode":"assert x.ndim > 0, f'unstack needs rank>0, got {x.shape}'","typeGuard":"def unstackable(x) -> bool:\n    return getattr(x, 'ndim', 0) > 0","tryCatchPattern":null,"preventionTips":["Keep dims with keepdims=True instead of over-squeezing","Check .ndim before unstack in generic utilities"],"tags":["jax","unstack","scalar","rank-error"],"backgroundTag":"rank-zero-array-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}