{"record":{"id":"375a85af7bdf255d","repo":"jax-ml/jax","slug":"unsupported-shape-shape","errorCode":null,"errorMessage":"Unsupported shape: {shape}","messagePattern":"Unsupported shape: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/_lowerings.py","lineNumber":143,"sourceCode":"\n  return core.ShapedArray(\n    shape=(shape[1] if transpose else shape[0], x.shape[1]),\n    dtype=x.dtype)\n\ndef _coo_spmm_gpu_lowering(ctx, data, row, col, x, *, transpose, shape,\n                           target_name_prefix):\n  data_aval, row_aval, _, x_aval = ctx.avals_in\n  nnz, = data_aval.shape\n  _, Ccols = x_aval.shape\n\n  batch_count = 1\n  if len(shape) == 2:\n    rows, cols = shape\n  elif len(shape) == 3:\n    batch_count, rows, cols = shape\n    nnz = nnz // batch_count\n  else:\n    raise NotImplementedError(f\"Unsupported shape: {shape}\")\n\n  # TODO(tianjianlu): use batch stride to trigger different mode of batch\n  # computation. Currently batch_stride = 0 is not allowed because of the issue\n  # in cusparse https://github.com/NVIDIA/CUDALibrarySamples/issues/81#issuecomment-1205562643\n  # Set batch stride to be the matrix size for now.\n  lhs_batch_stride = nnz\n  B_rows = rows if transpose else cols\n  rhs_batch_stride =  B_rows * Ccols\n\n  buffer_size, opaque = _get_module(target_name_prefix).build_coo_matmat_descriptor(\n      data_aval.dtype, x_aval.dtype, data_aval.dtype, row_aval.dtype,\n      rows, cols, Ccols, nnz, transpose, batch_count, lhs_batch_stride,\n      rhs_batch_stride)\n\n  buffer_aval = core.ShapedArray(shape=(buffer_size,), dtype=np.int8)\n  sub_ctx = ctx.replace(avals_out=[ctx.avals_out[0], buffer_aval])\n  rule = ffi.ffi_lowering(f\"{target_name_prefix}sparse_coo_matmat_ffi\")\n  return rule(sub_ctx, data, row, col, x, opaque=opaque)[:1]","sourceCodeStart":125,"sourceCodeEnd":161,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/_lowerings.py#L125-L161","documentation":"The COO sparse-matrix-times-dense GPU lowering only supports 2D matrices and 3D batched matrices. If the lhs shape has any other rank, the lowering cannot map it to a cuSPARSE routine and raises NotImplementedError.","triggerScenarios":"Calling matmul of a jax.experimental.sparse.COO with a dense array on GPU where the COO has rank other than 2 or 3 (e.g. extra dense dimensions making it 4D), via jnp.matmul / sp.dot dispatching to _coo_matmat_gpu_lowering.","commonSituations":"Using batched or higher-dimensional sparse arrays with COO format on GPU; COO lacks the dense-dimension and flexible batching support that BCOO has.","solutions":["Convert to BCOO (sparsify with BCOO or arr.tobcoo()) which supports arbitrary batch/dense dims on GPU","Reshape so the sparse operand is exactly 2D (or 3D batched) before the matmul","Run on CPU, where the COO CPU lowering supports more shapes"],"exampleFix":"// before\nout = coo_mat @ dense  # coo_mat has shape (B, N, M, K)\n// after\nbcoo_mat = coo_mat.reshape(...).tobcoo() if hasattr(coo_mat,'tobcoo') else coo_mat\nout = bcoo_mat @ dense","handlingStrategy":"fallback","validationCode":"shape = coo.shape\nif len(shape) not in (2, 3) and jax.devices()[0].platform == 'gpu':\n    coo = coo.reshape(...)  # or convert to BCOO","typeGuard":null,"tryCatchPattern":"try:\n    out = coo @ dense\nexcept NotImplementedError:\n    out = coo.tobcoo() @ dense","preventionTips":["Prefer BCOO over COO for anything beyond plain 2D on GPU","Check mat.ndim before GPU matmul with COO"],"tags":["jax","sparse","coo","gpu","shape-validation"],"backgroundTag":"unsupported-shape-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}