{"record":{"id":"f9669dbc9b8213b6","repo":"jax-ml/jax","slug":"a-and-b-must-be-2d-but-got-a-type-shape-and-b","errorCode":null,"errorMessage":"a and b must be 2D, but got: {a_type.shape} and {b_type.shape}","messagePattern":"a and b must be 2D, but got: (.+?) and (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2403,"sourceCode":"  elif isinstance(precision, tuple):\n    a_precision, b_precision = precision\n    if a_precision in _TF32_PRECISIONS or b_precision in _TF32_PRECISIONS:\n      input_precision = tt_dialect.InputPrecision.TF32\n    elif a_aval.dtype == jnp.float32:\n      input_precision = tt_dialect.InputPrecision.IEEE\n    else:\n      input_precision = None\n\n    acc_dtype = out_aval.dtype\n    if acc_dtype not in (jnp.int32, jnp.float16, jnp.float64):\n      acc_dtype = jnp.float32\n  else:\n    raise NotImplementedError(f\"Unsupported dot precision: {precision}.\")\n\n  a_type = ir.RankedTensorType(a.type)\n  b_type = ir.RankedTensorType(b.type)\n  if len(a_type.shape) != 2 or len(b_type.shape) != 2:\n    raise ValueError(\"a and b must be 2D, but got:\"\n                     f\" {a_type.shape} and {b_type.shape}\")\n\n  m, k = a_type.shape\n  _, n = b_type.shape\n  if a_type.element_type == ir.F64Type.get():\n    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates\n    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16\n    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK\n    # to round to zero, corrupting the ValueTable and segfaulting the compiler.\n    #   M >= 16  (2 × instrM=8)\n    #   N >=  8  (1 × instrN=8)\n    #   K >= 16  (4 × instrK=4)\n    errors = []\n    if m < 16:\n      errors.append(f\"M={m} < 16\")\n    if n < 8:\n      errors.append(f\"N={n} < 8\")\n    if k < 16:","sourceCodeStart":2385,"sourceCodeEnd":2421,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2385-L2421","documentation":"This internal check in the Pallas Triton dot lowering requires both operands to be rank-2 (2D) tensors when it decomposes shapes into (m, k) and (k, n). Higher-rank batched dots are not handled here and raise ValueError with the offending shapes.","triggerScenarios":"Passing 3D/4D arrays (e.g. (batch, m, k)) directly to the pallas dot lowering without first reshaping/batching manually; using pl.dot inside a kernel whose operands are 1D vectors.","commonSituations":"Writing a batched matmul kernel by reusing lax.dot_general semantics instead of looping over the grid; forgetting that Mosaic kernels require explicit 2D block operands.","solutions":["Reshape operands to 2D before the dot: a.reshape(m, k) @ block-level, or vmap/loop over batch via the kernel grid","Use jnp.dot or lax.dot_general outside pallas for batched matmuls","Ensure block specs give each dot exactly 2D tiles"],"exampleFix":"# before\nacc = pl.dot(a3d, b3d)  # shapes (B, M, K), (B, K, N)\n\n# after\n# launch grid with an extra dim and index 2D tiles\nacc = pl.dot(a3d[i], b3d[i])","handlingStrategy":"validation","validationCode":"def is_2d(a) -> bool:\n    return getattr(a, 'ndim', 0) == 2\nassert is_2d(a) and is_2d(b), f'{a.shape} {b.shape}'","typeGuard":"def is_2d_tile(a) -> bool:\n    return a.ndim == 2","tryCatchPattern":null,"preventionTips":["Reshape inputs to 2D blocks in kernel setup; use grid dims instead of batched dots"],"tags":["jax","pallas","triton","shape","matmul"],"backgroundTag":"rank-mismatch-in-matmul","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}