Function conv_dilate_window

axlearn/common/convolution.py:101–114

Returns dilated effective window size. Args: window: convolution window. dilation: convolution dilation. Returns: The dilated effective window size.

(*, window: Sequence[int], dilation: Optional[Sequence[int]])

# Copied from jax.lax._dilate_shape
# https://github.com/jax-ml/jax/blob/2d78b172266870bd755b039f6faa2056a51930f9/jax/_src/lax/lax.py#L5763
def conv_dilate_window(*, window: Sequence[int], dilation: Optional[Sequence[int]]):
    """Returns dilated effective window size.

    Args:
        window: convolution window.
        dilation: convolution dilation.

    Returns:
        The dilated effective window size.
    """
    if dilation is None or all(d == 1 for d in dilation):
        return window

    return tuple(max(1 + d * (w - 1), 0) for w, d in zip(window, dilation))


# Copied from subroutine in jax.lax.reduce_window.