Function conv_explicit_padding

axlearn/common/convolution.py:119–217

Returns the explicit padding for "SAME", "VALID", and "CAUSAL" modes. Each mode follows the formulas below: * SAME: (pad_total//2, pad_total - pad_total//2) s.t. pad_total = window-1 * VALID: (0, 0) * CAUSAL: (window - stride, stride - 1) Note: In the above equation, `window` w

(
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Optional[Sequence[int]] = None,
)

# Copied from subroutine in jax.lax.reduce_window.
# Extend lax.padtype_to_pads for CAUSAL.
def conv_explicit_padding(
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Optional[Sequence[int]] = None,
) -> ConvPaddingType:
    """Returns the explicit padding for "SAME", "VALID", and "CAUSAL" modes.

    Each mode follows the formulas below:
    * SAME: (pad_total//2, pad_total - pad_total//2) s.t. pad_total = window-1
    * VALID: (0, 0)
    * CAUSAL: (window - stride, stride - 1)

    Note: In the above equation, `window` will be replaced with `dilate_window` when dilation > 1.
        dilate_window = (window - 1) * dilation + 1. Check conv_dilate_window()

    For example, window=5 and stride=2,
    * SAME: padding = (2, 2)
                pad|           |pad
    paddings:   0 0|0 0 0 0 1 1|1 1
                |___^___|
                    |___^___|
                        |___^___|

    * VALID: padding = (0, 0)
               |           |
    paddings:  |0 0 0 0 1 1|
               |^_______|

    * CAUSAL: padding = (3, 1)
                pad  |           |pad
    paddings:   0 0 0|0 0 0 0 1 1|1
                |_____^_|
                    |_____^_|
                        |_____^_|


    For example, window=5, stride=2 and dilation=2
        -> dilate_window = 9 (== (window-1)*dilation + 1) and pad_total = 8
    * SAME: padding = (4, 4)
                    pad|                   |pad
    paddings:   0 0 0 0|0 0 0 0 0 0 0 0 1 1|1 1 1 1
                |_______^_______|
                    |_______^_______|
                        |_______^_______|
                            |_______^_______|
                                |_______^_______|

    * VALID: padding = (0, 0)
                |                   |pad
    paddings:   |0 0 0 0 0 0 0 0 1 1|
                |^_______________|

    * CAUSAL: padding = (7, 1)
                        pad  |                   |pad
    paddings:   0 0 0 0 0 0 0|0 0 0 0 0 0 0 0 1 1|1
                |_____________^_|