{"record":{"id":"50e997caefbc2db9","repo":"jax-ml/jax","slug":"matmul-with-object-of-shape-other-shape","errorCode":null,"errorMessage":"matmul with object of shape {other.shape}","messagePattern":"matmul with object of shape (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/coo.py","lineNumber":179,"sourceCode":"    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:\n      raise ValueError(f\"COO.tree_unflatten: invalid {aux_data=}\")\n    obj.shape = aux_data['shape']\n    obj._rows_sorted = aux_data['rows_sorted']\n    obj._cols_sorted = aux_data['cols_sorted']\n    return obj\n\n  def __matmul__(self, other: ArrayLike) -> Array:\n    if isinstance(other, JAXSparse):\n      raise NotImplementedError(\"matmul between two sparse objects.\")\n    other = jnp.asarray(other)\n    data, other = promote_dtypes(self.data, other)\n    self_promoted = COO((data, self.row, self.col), **self._info._asdict())\n    if other.ndim == 1:\n      return coo_matvec(self_promoted, other)\n    elif other.ndim == 2:\n      return coo_matmat(self_promoted, other)\n    else:\n      raise NotImplementedError(f\"matmul with object of shape {other.shape}\")\n\n#--------------------------------------------------------------------\n# coo_todense\n\ncoo_todense_p = core.Primitive('coo_todense')\n\ndef coo_todense(mat: COO) -> Array:\n  \"\"\"Convert a COO-format sparse matrix to a dense matrix.\n\n  Args:\n    mat : COO matrix\n  Returns:\n    mat_dense: dense version of ``mat``\n  \"\"\"\n  return _coo_todense(mat.data, mat.row, mat.col, spinfo=mat._info)\n\ndef _coo_todense(data: Array, row: Array, col: Array, *, spinfo: COOInfo) -> Array:\n  \"\"\"Convert CSR-format sparse matrix to a dense matrix.","sourceCodeStart":161,"sourceCodeEnd":197,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/coo.py#L161-L197","documentation":"COO.__matmul__ only accepts dense right operands with ndim 1 (matvec) or 2 (matmat). A dense operand with 3+ dimensions (or an otherwise odd ndim) hits the final NotImplementedError. Reduce/reshape the operand, or use BCOO's batched matmul.","triggerScenarios":"coo_mat @ dense_3d_tensor, or '@' with an operand whose ndim is not 1 or 2 (e.g. a batched stack of vectors with shape (B, N, 1)).","commonSituations":"Feeding a batched weight tensor directly instead of looping/vmapping; leftover extra size-1 dims from broadcasting in a dense pipeline (shape (N, 1) vs (N,)).","solutions":["Squeeze/reshape the operand to ndim 1 or 2 (e.g. other.reshape(other.shape[-2:]) per batch or other.squeeze())","Batch the product with jax.vmap over leading dims","Use BCOO and bcoo.bcoo_matmul / sparse.sparsify for natively batched sparse-dense products"],"exampleFix":"# before\nc = coo_mat @ W  # W.shape == (4, 8, 8) -> NotImplementedError\n\n# after\nc = jax.vmap(lambda w: coo_mat @ w)(W)","handlingStrategy":"validation","validationCode":"assert jnp.asarray(other).ndim in (1, 2), f'ndim={other.ndim} unsupported'","typeGuard":"def coo_matmul_shape_ok(other) -> bool:\n    import jax.numpy as jnp\n    return jnp.asarray(other).ndim in (1, 2)","tryCatchPattern":"try:\n    c = coo_mat @ W\nexcept NotImplementedError:\n    c = jax.vmap(lambda w: coo_mat @ w)(W)","preventionTips":["Squeeze size-1 dims before matmul","vmap over leading batch dims of the dense operand"],"tags":["jax","sparse","coo","matmul","shape-validation"],"backgroundTag":"sparse-matmul-dim-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}