{"record":{"id":"5c4f8a00321f4189","repo":"jax-ml/jax","slug":"collapseshapeop-with-empty-reassociation-is-not-su","errorCode":null,"errorMessage":"CollapseShapeOp with empty reassociation is not supported.","messagePattern":"CollapseShapeOp with empty reassociation is not supported\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2056,"sourceCode":"      cs.Equals(source_var, dest_var),\n  ]\n  return cs.ConstraintSystem(constraints=constraints), {\n      source_var: [source],\n      dest_var: [dest],\n  }\n\n\n@_add_constraint_system_derivation_rule(memref.CollapseShapeOp)\ndef _memref_collapse_shape_op_constraint_system(\n    ctx: DerivationContext,\n    op: memref.CollapseShapeOp,\n) -> ConstraintSystemDerivationRuleResult:\n  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)\n  # This should only occur when going from a (1, ...) shape to an empty shape.\n  # We can handle it if needed, but right now `CollapseShape` will not deal with\n  # this case.\n  if not reassociation:\n    raise NotImplementedError(\n        \"CollapseShapeOp with empty reassociation is not supported.\"\n    )\n\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  source_var = ctx.producer_ref(source)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  dest_var = cs.Variable(dest)\n\n  strides, _ = ir.MemRefType(source.value.type).get_strides_and_offset()\n  # In this case, we'd need additional checks to produce a correct constraint.\n  if strides != utils.get_contiguous_strides(source.shape):\n    raise NotImplementedError(\n        \"CollapseShapeOp with non-contiguous strides is not supported.\"\n    )\n\n  # TODO(bchetioui): We could generate an inverse expression `ExpandShape` in\n  # order to allow inferring layouts bidirectionally. This would allow removing\n  # transforms from some kernels' BlockSpecs, but is not necessary at this time.","sourceCodeStart":2038,"sourceCodeEnd":2074,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2038-L2074","documentation":"Mosaic GPU's layout inference cannot derive constraints for memref.CollapseShapeOp when the reassociation list is empty. An empty reassociation only occurs when collapsing a (1, ...) shape down to an empty (scalar-like) shape, a case the pass explicitly declines to handle.","triggerScenarios":"Calling memref.collapse_shape with reassociation=() on a memref whose shape is all 1s (e.g. (1,1) -> ()), producing an empty reassociation attribute inside a Mosaic GPU kernel.","commonSituations":"Generic reshape utilities that collapse all singleton dimensions; kernels parameterized so a tile shape degenerates to (1,...,1) at small sizes, triggering the empty-reassociation path.","solutions":["Guard your code path: skip the collapse (or use the original memref) when the reassociation would be empty, e.g. when all dimensions are size 1","Instead of collapsing to an empty shape, collapse to shape [1] with reassociation [[0,...]] so reassociation is non-empty","Adjust tile sizes/block specs so shapes never degenerate to all-ones"],"exampleFix":"# before\ncollapsed = memref.collapse_shape(src, reassociation=[])\n\n# after\nif src.shape and all(d == 1 for d in src.shape):\n    collapsed = memref.collapse_shape(src, reassociation=[[0] * len(src.shape)])\nelse:\n    collapsed = memref.collapse_shape(src, reassociation=reassoc)","handlingStrategy":"type-guard","validationCode":"reassoc = build_reassociation(src.shape, dest.shape)\nif not reassoc:  # collapsing (1,...) to ()\n    skip_collapse = True","typeGuard":"def is_collapse_supported(reassociation) -> bool:\n    return len(reassociation) > 0","tryCatchPattern":null,"preventionTips":["Never collapse to an empty/scalar shape; keep at least one dim of size 1","Guard shape-specialized kernels so all-ones tile shapes take a scalar code path"],"tags":["jax","mosaic-gpu","memref","collapse-shape","layout-inference","not-implemented"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}