{"record":{"id":"907def87e0c8e505","repo":"jax-ml/jax","slug":"arguments-to-jax-scipy-linalg-block-diag-must-have","errorCode":null,"errorMessage":"Arguments to jax.scipy.linalg.block_diag must have at most 2 dimensions, got {} at argument {}.","messagePattern":"Arguments to jax\\.scipy\\.linalg\\.block_diag must have at most 2 dimensions, got (.+?) at argument (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1658,"sourceCode":"\n  Examples:\n    >>> A = jnp.ones((1, 1))\n    >>> B = jnp.ones((2, 2))\n    >>> C = jnp.ones((3, 3))\n    >>> jax.scipy.linalg.block_diag(A, B, C)\n    Array([[1., 0., 0., 0., 0., 0.],\n           [0., 1., 1., 0., 0., 0.],\n           [0., 1., 1., 0., 0., 0.],\n           [0., 0., 0., 1., 1., 1.],\n           [0., 0., 0., 1., 1., 1.],\n           [0., 0., 0., 1., 1., 1.]], dtype=float32)\n  \"\"\"\n  if len(arrs) == 0:\n    arrs = (jnp.zeros((1, 0)),)\n  arrs = tuple(promote_dtypes(*arrs))\n  bad_shapes = [i for i, a in enumerate(arrs) if np.ndim(a) > 2]\n  if bad_shapes:\n    raise ValueError(\"Arguments to jax.scipy.linalg.block_diag must have at \"\n                     \"most 2 dimensions, got {} at argument {}.\"\n                     .format(arrs[bad_shapes[0]], bad_shapes[0]))\n  converted_arrs = [jnp.atleast_2d(a) for a in arrs]\n  dtype = lax.dtype(converted_arrs[0])\n  total_cols = sum(a.shape[1] for a in converted_arrs)\n\n  padded_arrs = []\n  current_col = 0\n  for arr in converted_arrs:\n    cols = arr.shape[1]\n    padding_config = ((0, 0, 0), (current_col, total_cols - cols - current_col, 0))\n    padded_arrs.append(lax.pad(arr, dtype.type(0), padding_config))\n    current_col += cols\n  return jnp.concatenate(padded_arrs, axis=0)\n\n\n@jit(static_argnames=(\"eigvals_only\", \"select\", \"select_range\"))\ndef eigh_tridiagonal(d: ArrayLike, e: ArrayLike, *, eigvals_only: bool = False,","sourceCodeStart":1640,"sourceCodeEnd":1676,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1640-L1676","documentation":"jax.scipy.linalg.block_diag builds a block-diagonal matrix from its arguments and requires each to be at most 2-D (scalars, vectors, or matrices). Any argument with ndim > 2 raises ValueError, reporting the offending array and its position in the argument list.","triggerScenarios":"Calling block_diag(a, b) where one argument is 3-D, e.g. a stacked batch of shape (4, 2, 2) from vmap or stacking.","commonSituations":"Feeding an accidentally stacked/extra-dimension array (e.g. result of jnp.stack or an un-squeezed matrix) into block_diag; SciPy has the same 2-D limit, so ported code with shape bugs surfaces here.","solutions":["Squeeze/reshape the offending argument to 2-D, e.g. arr = arr.reshape(-1, arr.shape[-1]) or arr[0] if batched","Check the argument index reported in the message to find which input is over-dimensional"],"exampleFix":"// before\nblocks = jnp.stack([m1, m2])  # (2, 2, 2)\nbd = jax.scipy.linalg.block_diag(blocks, m3)\n// after\nbd = jax.scipy.linalg.block_diag(m1, m2, m3)","handlingStrategy":"validation","validationCode":"bad = [i for i, a in enumerate(arrs) if np.ndim(a) > 2]\nassert not bad, f'args {bad} exceed 2 dims'","typeGuard":"def all_at_most_2d(arrs): return all(np.ndim(a) <= 2 for a in arrs)","tryCatchPattern":null,"preventionTips":["Unstack batches before block_diag; vmap block_diag if you need per-sample block diagonals","Check ndim of stacked intermediates (result of jnp.stack) before passing"],"tags":["jax","linalg","block-diagonal","shape-mismatch"],"backgroundTag":"invalid-array-dimensions","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}