Function pool_2d 

pub fn pool_2d(
    input: &[f32],
    input_mask: &[u8],
    input_shape: &[usize],
    config: &PoolConfig,
) -> (Vec<f32>, Vec<u8>, Vec<usize>)

Apply 2D spatial pooling to a gathered observation buffer.

input is the flat gather buffer in row-major order. input_mask has 1 for valid cells, 0 for padding. input_shape is [H, W].

Returns (output, output_mask, output_shape) where output_shape is [(H - kernel_size) / stride + 1, (W - kernel_size) / stride + 1].