For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Python class

Conv2d

`Conv2d`

class max.experimental.nn.Conv2d(kernel_size, in_channels, out_channels, dtype=None, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)

source

Bases: Module

A 2D convolution layer.

This is a Conv2d implementation that uses Tensor instead of Weight objects.

Example:

from max.driver import Accelerator
from max.dtype import DType
from max.experimental.nn import Conv2d
from max.experimental.tensor import Tensor

conv = Conv2d(
    kernel_size=3,
    in_channels=3,
    out_channels=64,
    dtype=DType.float32,
    has_bias=True,
    permute=True,
)
x = Tensor.ones([1, 3, 32, 32], dtype=DType.float32)
result = conv(x.to(Accelerator()))

Initialize Conv2d layer.

Parameters:

kernel_size (int | tuple[int, int]) – Size of the convolving kernel. Can be a single int (square kernel) or tuple (height, width).
in_channels (int) – Number of channels in the input image.
out_channels (int) – Number of channels produced by the convolution.
dtype (DType | None) – The data type for both weights and bias. In v3, this is optional as Tensor manages dtype automatically.
stride (int | tuple[int, int]) – Stride of the convolution for height and width dimensions. Can be int (applied to both dimensions) or tuple (stride_h, stride_w). Default: 1
padding (int | tuple[int, int] | tuple[int, int, int, int]) – Padding added to input. Can be int (applied to all sides), tuple of 2 ints (pad_h, pad_w), or tuple of 4 ints (pad_top, pad_bottom, pad_left, pad_right) to support asymmetric padding. Default: 0
dilation (int | tuple[int, int]) – Spacing between kernel elements for height and width dimensions. Can be int (applied to both dimensions) or tuple (dilation_h, dilation_w). Default: 1
num_groups (int) – Number of blocked connections from input channels to output channels. Input channels and output channels are divided into groups. Default: 1
device (Device | DeviceRef | None) – The target device for computation. In v3, this is optional as Tensor manages device automatically.
has_bias (bool) – If true, adds a learnable bias vector to the layer. Defaults to False.
permute (bool) – If true, permutes weights from PyTorch format to MAX format. PyTorch order: (out_channels, in_channels / num_groups, height, width). MAX API order: (height, width, in_channels / num_groups, out_channels). Defaults to False.
name (str | None) – Base name for weights. In v3, this is stored but not used for Weight naming. Defaults to None.

`bias`

bias: Tensor | Literal[0]

source

The bias tensor with shape [out_channels] (or 0 if bias is disabled).

`forward()`

forward(x)

source

Applies the 2D convolution to x.

Parameters:: x (Tensor) – The input tensor, shaped [batch_size, in_channels, height, width] when permute is True and [batch_size, height, width, in_channels] when permute is False.
Returns:: The output tensor, shaped [batch_size, out_channels, new_height, new_width] when permute is True and [batch_size, new_height, new_width, out_channels] when permute is False.
Return type:: Tensor

`weight`

weight: Tensor

source

The weight tensor with shape [out_channels, in_channels // num_groups, kernel_height, kernel_width].

Conv2d​

bias​

forward()​

weight​

`Conv2d`

`bias`

`forward()`

`weight`