{"record":{"id":"fc989530c6e49a6a","repo":"jax-ml/jax","slug":"mismatched-list-depths-in-jax-numpy-block","errorCode":null,"errorMessage":"Mismatched list depths in jax.numpy.block","messagePattern":"Mismatched list depths in jax\\.numpy\\.block","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4999,"sourceCode":"\n  arr, *choices = broadcast_arrays(arr, *choices)\n  return array(choices)[(arr,) + indices(arr.shape, sparse=True)]\n\n\ndef _atleast_nd(x: ArrayLike, n: int) -> Array:\n  m = np.ndim(x)\n  return lax.broadcast(x, (1,) * (n - m)) if m < n else asarray(x)\n\ndef _block(xs: ArrayLike | list[Any]) -> tuple[Array, int]:\n  if isinstance(xs, tuple):\n    raise ValueError(\"jax.numpy.block does not allow tuples, got {}\"\n                     .format(xs))\n  elif isinstance(xs, list):\n    if len(xs) == 0:\n      raise ValueError(\"jax.numpy.block does not allow empty list arguments\")\n    xs_tup, depths = unzip2([_block(x) for x in xs])\n    if any(d != depths[0] for d in depths[1:]):\n      raise ValueError(\"Mismatched list depths in jax.numpy.block\")\n    rank = max(depths[0], max(np.ndim(x) for x in xs_tup))\n    xs_tup = tuple(_atleast_nd(x, rank) for x in xs_tup)\n    return concatenate(xs_tup, axis=-depths[0]), depths[0] + 1\n  else:\n    return asarray(xs), 1\n\n\n@export\n@api.jit\ndef block(arrays: ArrayLike | list[Any]) -> Array:\n  \"\"\"Create an array from a list of blocks.\n\n  JAX implementation of :func:`numpy.block`.\n\n  Args:\n    arrays: an array, or nested list of arrays which will be concatenated\n      together to form the final array.\n","sourceCodeStart":4981,"sourceCodeEnd":5017,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L4981-L5017","documentation":"jnp.block requires the nested list structure to be a proper rectangular grid: every sub-list at the same level must have the same nesting depth. _block computes depths recursively and raises this ValueError when sibling depths differ, e.g. mixing a scalar with a list-of-lists.","triggerScenarios":"jnp.block([[a, b], c]) — second row is a raw array while the first is a list of two; or jnp.block([a, [b, c]]) where one element is a leaf and the other a list.","commonSituations":"Heterogeneous construction where some rows are single blocks (not wrapped in their own list) and others are multiple blocks; ported NumPy code that happened to work because shapes still concatenated correctly.","solutions":["Wrap scalar/single-array rows in their own list: [[a, b], [c]] becomes [[a, b], [c, ]]... ensure consistent nesting","Validate depths before calling: all(len(r) == len(rows[0]) and isinstance(r, list) for r in rows)","Assemble via explicit jnp.concatenate calls per axis instead of block"],"exampleFix":"// before\nout = jnp.block([[a, b], c])\n// after\nout = jnp.block([[a, b], [c]])","handlingStrategy":"validation","validationCode":"def uniform_depth(xs):\n    if not isinstance(xs, list):\n        return 0\n    depths = {uniform_depth(x) for x in xs}\n    return 1 + depths.pop() if len(depths) == 1 else -1\nassert uniform_depth(blocks) >= 0","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap every block entry in its own list for consistent nesting","Prefer explicit concatenate/stack when structure varies dynamically"],"tags":["jax","block","shape-mismatch","valueerror"],"backgroundTag":"ragged-nested-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}