{"record":{"id":"f1dd1f0e77a1e3a7","repo":"jax-ml/jax","slug":"matmul-with-object-of-shape-other-shape-f1dd1f","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/csr.py","lineNumber":134,"sourceCode":"\n  def todense(self):\n    return csr_todense(self)\n\n  def transpose(self, axes=None):\n    assert axes is None\n    return CSC((self.data, self.indices, self.indptr), shape=self.shape[::-1])\n\n  def __matmul__(self, other):\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    if other.ndim == 1:\n      return _csr_matvec(data, self.indices, self.indptr, other, shape=self.shape)\n    elif other.ndim == 2:\n      return _csr_matmat(data, self.indices, self.indptr, other, shape=self.shape)\n    else:\n      raise NotImplementedError(f\"matmul with object of shape {other.shape}\")\n\n  def tree_flatten(self):\n    return (self.data, self.indices, self.indptr), {\"shape\": self.shape}\n\n  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    obj = object.__new__(cls)\n    obj.data, obj.indices, obj.indptr = children\n    if aux_data.keys() != {'shape'}:\n      raise ValueError(f\"CSR.tree_unflatten: invalid {aux_data=}\")\n    obj.__dict__.update(**aux_data)\n    return obj\n\n\n@tree_util.register_pytree_node_class\nclass CSC(JAXSparse):\n  \"\"\"Experimental CSC matrix implemented in JAX; API subject to change.\"\"\"\n  data: jax.Array","sourceCodeStart":116,"sourceCodeEnd":152,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/csr.py#L116-L152","documentation":"CSR.__matmul__ dispatches on the dense right operand's ndim: 1 → _csr_matvec, 2 → _csr_matmat; anything else (3+ dims, or 0-d) falls through to NotImplementedError with the offending shape. Batch the product or use BCOO/BCSR which support batched operands.","triggerScenarios":"csr_mat @ tensor3d where tensor3d.ndim >= 3 (e.g. shape (B, N, M)), or '@' with a 0-d array.","commonSituations":"Passing a batched stack of matrices/vectors from a dense pipeline directly; leftover size-1 dims making a vector ndim 3; vectorized training loops that assume broadcasting like numpy matmul.","solutions":["Use jax.vmap over the leading batch dims: jax.vmap(lambda w: csr @ w)(tensor)","Reshape/squeeze the operand to ndim 1 or 2","Switch to BCOO (bcoo_matmul) or BCSR for natively batched sparse-dense matmul"],"exampleFix":"# before\nc = csr_mat @ W  # W.shape == (4, 8, 8) -> NotImplementedError\n\n# after\nc = jax.vmap(lambda w: csr_mat @ w)(W)","handlingStrategy":"validation","validationCode":"assert jnp.asarray(other).ndim in (1, 2), f'ndim={other.ndim} unsupported by CSR matmul'","typeGuard":"def csr_matmul_shape_ok(other) -> bool:\n    import jax.numpy as jnp\n    return jnp.asarray(other).ndim in (1, 2)","tryCatchPattern":"try:\n    c = csr_mat @ W\nexcept NotImplementedError:\n    c = jax.vmap(lambda w: csr_mat @ w)(W)","preventionTips":["vmap over batch dims of dense operands","Squeeze stray size-1 dims before '@'"],"tags":["jax","sparse","csr","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"}