code: def _pooling_general(inputs, reducer, init_val, rescaler=None,
                       pool_size=(2, 2), strides=None, padding='VALID'):
    
    spatial_strides = strides or (1,) * len(pool_size)
    rescale = rescaler(pool_size, spatial_strides, padding) if rescaler else None
    dims = (1,) + pool_size + (1,)  # NHWC
    strides = (1,) + spatial_strides + (1,)
    out = lax.reduce_window(inputs, init_val, reducer, dims, strides, padding)
    return rescale(out, inputs) if rescale else out
docs: Helper: general pooling computation used in pooling layers later.