IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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.experimental.nn import Conv2d
from max.experimental.tensor import Tensor

conv = Conv2d(
    kernel_size=3,
    in_channels=3,
    out_channels=64,
    has_bias=True,
    permute=True,
)

x = Tensor.ones([1, 3, 32, 32])
result = conv(x)

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

Apply 2D convolution to input.

Parameters:

x (Tensor) – Input tensor. Shape depends on permute:

  • If permute=True: [batch_size, in_channels, height, width]
  • If permute=False: [batch_size, height, width, in_channels]

Returns:

  • If permute=True: [batch_size, out_channels, new_height, new_width]
  • If permute=False: [batch_size, new_height, new_width, out_channels]

Return type:

Output tensor. Shape depends on permute

weight​

weight: Tensor

source

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