Function _emit_pool_nd

backends/mlx/ops.py:3522–3693

Emit IR nodes for N-dimensional pooling. Decomposes pooling into: Transpose (channels-first -> channels-last) -> Pad (if needed) -> Reshape+Transpose (fast path) or AsStrided (general path) -> Max/Mean reduction over kernel dims -> Transpose (channels-last -> chann

(
    P: MLXProgramBuilder,
    n: Node,
    ndim: int,
    reduce_node_cls: type,
    padding_value: float,
    kernel_size: List[int],
    stride: List[int],
    padding: List[int],
)

def _emit_pool_nd(
    P: MLXProgramBuilder,
    n: Node,
    ndim: int,
    reduce_node_cls: type,
    padding_value: float,
    kernel_size: List[int],
    stride: List[int],
    padding: List[int],
) -> Slot:
    """Emit IR nodes for N-dimensional pooling.

    Decomposes pooling into:
      Transpose (channels-first -> channels-last)
      -> Pad (if needed)
      -> Reshape+Transpose (fast path) or AsStrided (general path)
      -> Max/Mean reduction over kernel dims
      -> Transpose (channels-last -> channels-first)

    Works for 1D, 2D, and 3D pooling uniformly.

    Args:
        P: Program builder.
        n: FX graph node for the pooling op.
        ndim: Spatial dimensionality (1, 2, or 3).
        reduce_node_cls: MaxNode or MeanNode.
        padding_value: Padding fill value (-inf for max, 0 for avg).
        kernel_size: Kernel size per spatial dim, length ndim.
        stride: Stride per spatial dim, length ndim.
        padding: Padding per spatial dim, length ndim.

    Returns:
        Output Slot with shape [N, C, *out_spatial].
    """
    x_node = P.args(n)[0]
    (x,) = P.slot_map([x_node])
    x_meta = n.args[0].meta["val"]
    shape = list(x_meta.shape)  # [N, C, *spatial]

    N = shape[0]
    C = shape[1]
    spatial = shape[2:]  # length == ndim

    # 1. Transpose: channels-first [N, C, *spatial] -> channels-last [N, *spatial, C]
    to_cl = [0] + list(range(2, ndim + 2)) + [1]
    _, cur = P.make_tmp_slot()
    P.emit(
        TransposeNode(
            x=P.slot_to_tid(x),
            out=P.slot_to_tid(cur),
            perm=to_cl,
        )
    )

    # 2. Pad spatial dims if needed
    spatial_padded = [s + 2 * p for s, p in zip(spatial, padding)]
    if any(p > 0 for p in padding):
        pad_width = [0, 0]  # batch dim: no pad