keras-team/keras · error · ValueError

adaptive_average_pool supports only 1D/2D/3D inputs

Error message

adaptive_average_pool supports only 1D/2D/3D inputs

What it means

Dispatch guard in the JAX backend's adaptive_average_pool: it computes the number of spatial dims as inputs.ndim - 2 and routes 1D/2D/3D cases to dedicated kernels. Any other rank (0D, 4D+, or inputs with fewer than 3 dims) falls through to this generic rejection. Note the message text says 'adaptive_average_pool' but it applies to the average variant.

Solutions

1. Reshape the input to rank 3 (1D: (batch, length, channels)), rank 4 (2D), or rank 5 (3D) depending on the pooling variant used
2. For higher-dimensional pooling, apply the op to a flattened spatial view or use regular (non-adaptive) pooling layer-by-layer
3. Verify data_format matches the actual layout so ndim is interpreted correctly

When it happens

Trigger: Thrown at keras/src/backend/jax/ops/nn.py:684 when the library encounters an invalid state.

Common situations: See trigger scenarios.

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

Appendix: source

Thrown at keras/src/backend/jax/ops/nn.py:684

    combined_w = jnp.concatenate([small_w_pool, big_w_pool], axis=3)
    out = jnp.take(combined_w, gather_w, axis=3)

    if data_format == "channels_first":
        out = jnp.transpose(out, (0, 4, 1, 2, 3))

    return out


def adaptive_average_pool(inputs, output_size, data_format=None):
    data_format = standardize_data_format(data_format)
    dims = inputs.ndim - 2
    if dims == 1:
        return _adaptive_average_pool1d(inputs, output_size, data_format)
    if dims == 2:
        return _adaptive_average_pool2d(inputs, output_size, data_format)
    if dims == 3:
        return _adaptive_average_pool3d(inputs, output_size, data_format)
    raise ValueError("adaptive_average_pool supports only 1D/2D/3D inputs")


def adaptive_max_pool(inputs, output_size, data_format=None):
    data_format = standardize_data_format(data_format)
    dims = inputs.ndim - 2
    if dims == 1:
        return _adaptive_max_pool1d(inputs, output_size, data_format)
    if dims == 2:
        return _adaptive_max_pool2d(inputs, output_size, data_format)
    if dims == 3:
        return _adaptive_max_pool3d(inputs, output_size, data_format)
    raise ValueError("adaptive_max_pool supports only 1D/2D/3D inputs")


def _convert_to_lax_conv_dimension_numbers(
    num_spatial_dims,
    data_format="channels_last",
    transpose=False,

View on GitHub (pinned to 7a34a03db6)