{"record":{"id":"42f20729c974fa0f","repo":"jax-ml/jax","slug":"coo-must-have-ndim-2-got-shape","errorCode":null,"errorMessage":"COO must have ndim=2; got {shape=}","messagePattern":"COO must have ndim=2; got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/coo.py","lineNumber":119,"sourceCode":"    \"\"\"Return a copy of the COO matrix with sorted indices.\n\n    The matrix is sorted by row indices and column indices per row.\n    If self._rows_sorted is True, this returns ``self`` without a copy.\n    \"\"\"\n    # TODO(jakevdp): would be benefit from lowering this to cusparse sort_rows utility?\n    if self._rows_sorted:\n      return self\n    row, col, data = lax.sort((self.row, self.col, self.data), num_keys=2)\n    return self.__class__((data, row, col), shape=self.shape,\n                          rows_sorted=True)\n\n  @classmethod\n  def _empty(cls, shape: Sequence[int], *, dtype: DTypeLike | None = None,\n             index_dtype: DTypeLike = 'int32') -> COO:\n    \"\"\"Create an empty COO instance. Public method is sparse.empty().\"\"\"\n    shape = tuple(shape)\n    if len(shape) != 2:\n      raise ValueError(f\"COO must have ndim=2; got {shape=}\")\n    data = jnp.empty(0, dtype)\n    row = col = jnp.empty(0, index_dtype)\n    return cls((data, row, col), shape=shape, rows_sorted=True,\n               cols_sorted=True)\n\n  @classmethod\n  def _eye(cls, N: int, M: int, k: int, *, dtype: DTypeLike | None = None,\n           index_dtype: DTypeLike = 'int32') -> COO:\n    if k > 0:\n      diag_size = min(N, M - k)\n    else:\n      diag_size = min(N + k, M)\n\n    if diag_size <= 0:\n      # if k is out of range, return an empty matrix.\n      return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)\n\n    data = jnp.ones(diag_size, dtype=dtype)","sourceCodeStart":101,"sourceCodeEnd":137,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/coo.py#L101-L137","documentation":"The legacy jax.experimental.sparse.COO format only supports 2D matrices (row/col index buffers imply exactly two sparse dimensions). COO._empty, used by sparse.empty(format='coo') and sparse.eye(format='coo'), validates len(shape) == 2. For arbitrary-dimensional sparse arrays use BCOO.","triggerScenarios":"sparse.empty(shape, format='coo') or sparse.eye(N, format='coo') with a shape that is not length 2 — e.g. sparse.empty((2,3,4), format='coo') or sparse.eye on an ndim != 2 shape parameter.","commonSituations":"Older JAX code using COO being extended to batched/3D tensors; format strings chosen dynamically and hitting COO for non-matrix shapes.","solutions":["Use format='bcoo' (jax.experimental.sparse.BCOO) for any non-2D sparse array","Reshape/flatten the problem to 2D if the legacy COO API is required","Prefer BCOO in new code even for 2D — COO/CSR/CSC are legacy wrappers"],"exampleFix":"# before\nm = sparse.empty((2, 3, 4), format='coo')  # ValueError: ndim=2 required\n\n# after\nm = sparse.empty((2, 3, 4), format='bcoo')","handlingStrategy":"validation","validationCode":"assert len(tuple(shape)) == 2, 'COO is 2D only; use bcoo'","typeGuard":"def coo_shape_ok(shape) -> bool:\n    return len(tuple(shape)) == 2","tryCatchPattern":"try:\n    m = sparse.empty(shape, format='coo')\nexcept ValueError:\n    m = sparse.empty(shape, format='bcoo')","preventionTips":["Treat COO/CSR/CSC as 2D-only legacy formats","Default to bcoo in format-dispatch code"],"tags":["jax","sparse","coo","shape-validation"],"backgroundTag":"sparse-format-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}