Skip to content

Convolution

simplegrad.functions.conv.pad(x: Tensor, width: int | tuple[int, int, int, int], mode: str = 'constant', value: int = 0) -> Tensor

Pad a tensor along its spatial dimensions.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

required
width int | tuple[int, int, int, int]

Padding widths. An int or nested tuples as accepted by numpy.pad.

required
mode str

Padding mode (e.g. "constant", "reflect"). See numpy.pad docs.

'constant'
value int

Fill value for "constant" mode.

0

Returns:

Type Description
Tensor

Padded tensor.

Source code in simplegrad/functions/conv.py
def pad(
    x: Tensor, width: int | tuple[int, int, int, int], mode: str = "constant", value: int = 0
) -> Tensor:
    """Pad a tensor along its spatial dimensions.

    Args:
        x: Input tensor.
        width: Padding widths. An int or nested tuples as accepted by ``numpy.pad``.
        mode: Padding mode (e.g. ``"constant"``, ``"reflect"``). See numpy.pad docs.
        value: Fill value for ``"constant"`` mode.

    Returns:
        Padded tensor.
    """
    return _Pad.apply(x, width, mode, value, oper=f"pad(width={width}, mode={mode})")

simplegrad.functions.conv.conv2d(x: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | tuple[int, int] = 1, pad_width: int | tuple[int, int, int, int] = 0, pad_mode: str = 'constant', pad_value: int = 0) -> Tensor

Apply a 2D convolution over the input tensor.

Implements convolution as a single matrix multiplication (im2col). Supports batched inputs (4D) and single-sample inputs (3D).

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch, in_channels, H, W) or (in_channels, H, W).

required
weight Tensor

Kernel tensor of shape (out_channels, in_channels, kH, kW).

required
bias Tensor | None

Optional bias of shape (out_channels,).

None
stride int | tuple[int, int]

Convolution stride. Int or (height_stride, width_stride).

1
pad_width int | tuple[int, int, int, int]

Padding to apply before convolution. Int (same on all sides) or (top, bottom, left, right).

0
pad_mode str

Padding mode passed to numpy.pad. Defaults to "constant".

'constant'
pad_value int

Fill value for constant padding. Defaults to 0.

0

Returns:

Type Description
Tensor

Output tensor of shape (batch, out_channels, out_H, out_W).

Source code in simplegrad/functions/conv.py
def conv2d(
    x: Tensor,
    weight: Tensor,
    bias: Tensor | None = None,
    stride: int | tuple[int, int] = 1,
    # 2 options for padding:
    # int - same padding for all directions
    # tuple of 4 ints - (top, bottom, left, right)
    pad_width: int | tuple[int, int, int, int] = 0,
    pad_mode: str = "constant",
    pad_value: int = 0,
) -> Tensor:
    """Apply a 2D convolution over the input tensor.

    Implements convolution as a single matrix multiplication (im2col). Supports
    batched inputs (4D) and single-sample inputs (3D).

    Args:
        x: Input tensor of shape ``(batch, in_channels, H, W)`` or
            ``(in_channels, H, W)``.
        weight: Kernel tensor of shape ``(out_channels, in_channels, kH, kW)``.
        bias: Optional bias of shape ``(out_channels,)``.
        stride: Convolution stride. Int or ``(height_stride, width_stride)``.
        pad_width: Padding to apply before convolution. Int (same on all sides)
            or ``(top, bottom, left, right)``.
        pad_mode: Padding mode passed to ``numpy.pad``. Defaults to ``"constant"``.
        pad_value: Fill value for constant padding. Defaults to 0.

    Returns:
        Output tensor of shape ``(batch, out_channels, out_H, out_W)``.
    """
    assert (
        len(x.shape) == 4 or len(x.shape) == 3
    ), "Input tensor must be 4-dimensional (batch_size, in_channels, height, width) or 3-dimensional (in_channels, height, width)"
    assert len(weight.shape) == 4, "Weight tensor must be 4-dimensional"

    if pad_width == 0 or pad_width == (0, 0, 0, 0):
        padded_input = x
    else:
        if isinstance(pad_width, int):
            pad_width_np = (
                (0, 0),
                (0, 0),
                (pad_width, pad_width),
                (pad_width, pad_width),
            )
        elif isinstance(pad_width, tuple) and len(pad_width) == 4:
            pad_width_np = (
                (0, 0),
                (0, 0),
                (pad_width[0], pad_width[1]),
                (pad_width[2], pad_width[3]),
            )
        else:
            raise ValueError(
                "pad_width must be an int or a tuple of 4 ints (top, bottom, left, right)"
            )

        padded_input = pad(x=x, width=pad_width_np, mode=pad_mode, value=pad_value)

    return _conv2d_no_pad(padded_input=padded_input, weight=weight, bias=bias, stride=stride)