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

Capability guard in the Pallas Mosaic TPU conv lowering: feature_group_count or batch_group_count greater than 1 (grouped/depthwise convolution) is not implemented on this backend and raises NotImplementedError; use plain convolutions or reshape grouped convs manually.

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.

Appendix: 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)