keras-team/keras · error · ValueError

AWQ quantization only supports 2D or 3D kernels.

Error message

AWQ quantization only supports 2D or 3D kernels.

What it means

EinsumDense._awq_build only supports 2D and 3D kernels for AWQ quantization: group-wise scaling is defined over a rows-by-columns view of the kernel. A kernel of any other rank (e.g. 4D from a nested output_dim) makes the AWQ math undefined, so it raises this ValueError.

Source

Thrown at keras/src/layers/core/einsum_dense.py:817

            shape = list(self.original_kernel_shape)
            d_model_dim_index = shape.index(max(shape))

            if d_model_dim_index == 0:  # QKV projection case
                in_features, heads, head_dim = shape
                rows, columns = (
                    in_features,
                    heads * head_dim,
                )
            elif d_model_dim_index in [1, 2]:  # Attention Output case
                heads, head_dim, out_features = shape
                rows, columns = (
                    heads * head_dim,
                    out_features,
                )
            else:
                raise ValueError("Could not determine row/column split.")
        else:
            raise ValueError("AWQ quantization only supports 2D or 3D kernels.")

        group_size = awq_core.get_group_size_for_layer(self, config)
        num_groups = 1 if group_size == -1 else math.ceil(rows / group_size)

        self.awq_unpacked_column_size = columns

        # For 4-bit weights, we pack two values per byte.
        kernel_columns = (columns + 1) // 2

        self._set_quantization_info()

        self.quantized_kernel = self.add_weight(
            name="kernel",
            shape=(kernel_columns, rows),
            initializer="zeros",
            dtype="uint8",
            trainable=False,
        )

View on GitHub (pinned to 7a34a03db6)

Solutions

  1. Exclude this layer from AWQ quantization with filters.
  2. Reduce the kernel to rank 2 or 3 by flattening dimensions in the equation or splitting the layer into simpler EinsumDense ops.
  3. Pick a different quantization mode for this layer.

Example fix

# before
layer = keras.layers.EinsumDense('abcd,cde->abe', output_dim=(4, 8, 16))  # 4D kernel
model.quantize(awq_config)  # ValueError

# after
model.quantize(awq_config, filters=[l.name for l in model.layers if l is not layer])
Defensive patterns

Strategy: validation

Validate before calling

rank = len(tuple(layer.kernel.shape))
if rank not in (2, 3):
    raise RuntimeError(f'AWQ unsupported for rank-{rank} kernel: {layer.name}')

Type guard

def awq_compatible(layer) -> bool:
    return layer.built and len(tuple(layer.kernel.shape)) in (2, 3)

Prevention

When it happens

Trigger: model.quantize(...) with mode 'awq' including an EinsumDense whose output_dim/equation yields a 4D+ kernel (e.g. output_dim=(heads, head_dim, features, extra)).

Common situations: Fused or exotic projections with higher-rank kernels caught by a broad AWQ layer filter; models ported from other frameworks where EinsumDense is used as a general tensor contraction.

Related errors


AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/3bed9985b8edf376. Report an issue: GitHub.