Function _check_conv_cfg

axlearn/common/convolution.py:28–65

(
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Optional[Sequence[int]],
    input_dim: int,
    output_dim: int,
    num_input_dim_groups: int,
)

# TODO(yuanliu939): Make this take `BaseConv.Config` directly.
def _check_conv_cfg(
    *,
    window: Sequence[int],
    strides: Sequence[int],
    padding: ConvPaddingType,
    dilation: Optional[Sequence[int]],
    input_dim: int,
    output_dim: int,
    num_input_dim_groups: int,
):
    if any(w < 1 for w in window):
        raise ValueError(f"window ({window}) must be a positive integer.")

    if any(s < 1 for s in strides):
        raise ValueError(f"strides ({strides}) must be a positive integer.")

    if isinstance(padding, str):
        if padding not in SUPPORT_CONV_PADDING:
            raise ValueError(f"{padding} padding is not supported.")
    else:
        padding_flattened = jax.tree.leaves(padding)
        if any(p < 0 for p in padding_flattened):
            raise ValueError("Negative padding is not supported")

    if dilation is not None and any(d < 1 for d in dilation):
        raise ValueError(f"dilation ({dilation}) must be a positive integer.")

    if input_dim % num_input_dim_groups != 0:
        raise ValueError(
            f"input_dim ({input_dim}) must be divisible by "
            f"num_input_dim_groups({num_input_dim_groups})."
        )

    if output_dim % num_input_dim_groups != 0:
        raise ValueError(
            f"output_dim ({output_dim}) must be divisible by "
            f"num_input_dim_groups({num_input_dim_groups})."
        )


class BaseConv(BaseLayer):