{"record":{"id":"a0be652a9ac3d5b6","repo":"jax-ml/jax","slug":"grouped-convolutions-are-not-supported-on-pallas-m","errorCode":null,"errorMessage":"Grouped convolutions are not supported on Pallas Mosaic TPU backend yet.","messagePattern":"Grouped convolutions are not supported on Pallas Mosaic TPU backend yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3014,"sourceCode":"@register_lowering_rule(lax.conv_general_dilated_p)\ndef _conv_general_dilated_lowering_rule(\n    ctx: LoweringRuleContext,\n    lhs,\n    rhs,\n    *,\n    window_strides,\n    padding,\n    lhs_dilation,\n    rhs_dilation,\n    dimension_numbers,\n    feature_group_count,\n    batch_group_count,\n    precision=None,\n    preferred_element_type=None,\n    **_,\n):\n  if feature_group_count != 1 or batch_group_count != 1:\n    raise NotImplementedError(\n        \"Grouped convolutions are not supported on Pallas Mosaic TPU backend\"\n        \" yet.\"\n    )\n  return _conv_lowering_rule(\n      ctx,\n      lhs,\n      rhs,\n      dimension_numbers=dimension_numbers,\n      window_strides=window_strides,\n      padding=padding,\n      lhs_dilation=lhs_dilation,\n      rhs_dilation=rhs_dilation,\n      precision=precision,\n  )\n\n\n@register_lowering_rule(tpu_primitives.conv_p)\ndef _conv_lowering_rule(","sourceCodeStart":2996,"sourceCodeEnd":3032,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2996-L3032","documentation":"The grouped-convolution wrapper rejects feature_group_count != 1 or batch_group_count != 1 before delegating to the internal conv lowering, because Mosaic TPU has no grouped conv support yet.","triggerScenarios":"lax.conv_general_dilated with feature_group_count > 1 or batch_group_count > 1 inside a Pallas Mosaic TPU kernel (e.g. depthwise convolutions).","commonSituations":"Depthwise/grouped conv layers from CNN architectures (MobileNet-style) run through Pallas on TPU.","solutions":["Set feature_group_count=1 and batch_group_count=1; restructure depthwise conv as per-group standard convs or gather/scatter ops","Run grouped convs outside the Pallas kernel via regular JAX/XLA","Reshape so each group is handled by a separate kernel invocation or loop iteration"],"exampleFix":"// before\nout = lax.conv_general_dilated(x, w, ..., feature_group_count=C)\n// after\nouts = [lax.conv_general_dilated(x[:, g], w[g], ...) for g in range(C)]\nout = jnp.stack(outs, axis=1)","handlingStrategy":"validation","validationCode":"assert feature_group_count == 1 and batch_group_count == 1, 'grouped conv unsupported in Pallas Mosaic'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep grouped/depthwise convs outside Pallas kernels","Vmap or loop over groups instead of using group counts"],"tags":["jax","pallas","tpu","convolution","grouped-conv","depthwise"],"backgroundTag":"grouped-convolution-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}