jax-ml/jax · error · NotImplementedError

Grouped convolutions are not supported on Pallas Mosaic TPU

Error message

Grouped convolutions are not supported on Pallas Mosaic TPU backend yet.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3014

@register_lowering_rule(lax.conv_general_dilated_p)
def _conv_general_dilated_lowering_rule(
    ctx: LoweringRuleContext,
    lhs,
    rhs,
    *,
    window_strides,
    padding,
    lhs_dilation,
    rhs_dilation,
    dimension_numbers,
    feature_group_count,
    batch_group_count,
    precision=None,
    preferred_element_type=None,
    **_,
):
  if feature_group_count != 1 or batch_group_count != 1:
    raise NotImplementedError(
        "Grouped convolutions are not supported on Pallas Mosaic TPU backend"
        " yet."
    )
  return _conv_lowering_rule(
      ctx,
      lhs,
      rhs,
      dimension_numbers=dimension_numbers,
      window_strides=window_strides,
      padding=padding,
      lhs_dilation=lhs_dilation,
      rhs_dilation=rhs_dilation,
      precision=precision,
  )


@register_lowering_rule(tpu_primitives.conv_p)
def _conv_lowering_rule(

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = lax.conv_general_dilated(x, w, ..., feature_group_count=C)
// after
outs = [lax.conv_general_dilated(x[:, g], w[g], ...) for g in range(C)]
out = jnp.stack(outs, axis=1)

Defensive patterns

Strategy: validation

Validate before calling

assert feature_group_count == 1 and batch_group_count == 1, 'grouped conv unsupported in Pallas Mosaic'

Prevention

Keep grouped/depthwise convs outside Pallas kernels
Vmap or loop over groups instead of using group counts

When it happens

Trigger: lax.conv_general_dilated with feature_group_count > 1 or batch_group_count > 1 inside a Pallas Mosaic TPU kernel (e.g. depthwise convolutions).

Common situations: Depthwise/grouped conv layers from CNN architectures (MobileNet-style) run through Pallas on TPU.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/a0be652a9ac3d5b6. Report an issue: GitHub.