Skip to main content

Python module

conv

The conv module provides classes for performing convolution operations in various dimensions (1D, 2D, and 3D) on tensor inputs. These convolution operations are core building blocks for neural networks, especially in computer vision and sequence processing tasks.

Here’s an example demonstrating how to use a 1D convolution:

import max.nn as nn
from max.graph import Graph, ops, Weight, DeviceRef
from max.dtype import DType
import numpy as np

with Graph(name="conv_example") as graph:
    # Define dimensions
    batch_size = 2
    seq_length = 10
    in_channels = 16
    out_channels = 32
    kernel_size = 3

    # Create input tensor [batch_size, sequence_length, channels]
    x_data = np.zeros((batch_size, seq_length, in_channels), dtype=np.float32)
    x = ops.constant(x_data, dtype=DType.float32, device=DeviceRef.CPU())

    # Create weights for convolution
    filter_1d = Weight(
        name="filter_weight",
        dtype=DType.float32,
        shape=[kernel_size, in_channels, out_channels]
        device=DeviceRef.CPU()
    )
    bias_1d = Weight(
        name="bias_weight",
        dtype=DType.float32,
        shape=[out_channels]
        device=DeviceRef.CPU()
    )

    # Create and apply Conv1D layer
    conv1d = nn.Conv1D(
        filter=filter_1d,
        bias=bias_1d,
        stride=1,
        padding=1
    )

    output_1d = conv1d(x)
    print(f"Conv1D output shape: {output_1d.shape}")
    # Output: Conv1D output shape: [Dim(2), Dim(10), Dim(32)]

Conv1D

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

A 1D convolution over an input signal composed of several input planes.

Example:

conv = nn.Conv1D(
    kernel_size=3,
    in_channels=64,
    out_channels=128,
    dtype=DType.float32,
    stride=1,
    padding=0,
    has_bias=False,
    name="conv1d_weight",
    device=DeviceRef.GPU(),
)

Parameters:

bias

bias: Weight | None = None

The optional bias vector stored on CPU with shape (out_channels,). Model init moves the bias to device if present.

device

device: DeviceRef | None

The device where matrix operations are performed.

dilation

dilation: int

Controls the dilation rate.

filter

filter: Weight

The weight matrix stored on CPU with shape (kernel_size, in_channels / num_groups, out_channels). Model init moves the weight to device.

num_groups

num_groups: int

Number of blocked connections from input channels to output channels.

padding

padding: int | tuple[int, int]

Controls the amount of padding applied to the input.

If int: symmetric padding applied to both sides (pad_left = pad_right = padding). If tuple[int, int]: asymmetric padding as (pad_left, pad_right).

permute

permute: bool = False

bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, kernel_size) Max API order: (kernel_size, in_channels / num_groups, out_channels).

stride

stride: int

Controls the stride for the cross-correlation.

Conv2d

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

A 2D convolution over an input signal composed of several input planes.

Example:

conv = nn.Conv2d(
    kernel_size=3,
    in_channels=64,
    out_channels=128,
    dtype=DType.float32,
    stride=1,
    padding=0,
    has_bias=False,
    name="conv2d_weight",
    device=DeviceRef.GPU(),
)

Parameters:

bias

bias: Weight | None = None

The optional bias vector stored on CPU with shape (out_channels,). Model init moves the bias to device if present.

device

device: DeviceRef | None

The device where matrix operations are performed.

dilation

dilation: tuple[int, int]

Controls the dilation rate.

filter

filter: Weight

The weight matrix stored on CPU with shape (height, width, in_channels / num_groups, out_channels). Model init moves the weight to device.

num_groups

num_groups: int

Number of blocked connections from input channels to output channels.

padding

padding: tuple[int, int, int, int]

Controls the amount of padding applied before and after the input for height and width dimensions.

Format: (pad_top, pad_bottom, pad_left, pad_right).

permute

permute: bool = False

bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, height, width) Max API order: (height, width, in_channels / num_groups, out_channels).

shard()

shard(devices)

Creates sharded views of this Conv2d layer across multiple devices.

Parameters:

devices (Iterable[DeviceRef]) – Iterable of devices to place the shards on.

Returns:

List of sharded Conv2d instances, one for each device.

Return type:

list[Conv2d]

sharding_strategy

property sharding_strategy: ShardingStrategy | None

Get the Conv2d sharding strategy.

stride

stride: tuple[int, int]

Controls the stride for the cross-correlation.

Conv3D

class max.nn.legacy.conv.Conv3D(depth, height, width, in_channels, out_channels, dtype, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)

A 3D convolution over an input signal composed of several input planes.

Example:

conv = nn.Conv3D(
    depth=3,
    height=3,
    width=3,
    in_channels=64,
    out_channels=128,
    dtype=DType.float32,
    stride=1,
    padding=0,
    has_bias=False,
    name="conv3d_weight",
    device=DeviceRef.GPU(),
)

Parameters:

bias

bias: Weight | None = None

The optional bias vector stored on CPU with shape (out_channels,). Model init moves the bias to device if present.

device

device: DeviceRef | None

The device where matrix operations are performed.

dilation

dilation: tuple[int, int, int]

Controls the dilation rate for depth, height, and width dimensions.

filter

filter: Weight

The weight matrix stored on CPU with shape (depth, height, width, in_channels / num_groups, out_channels). Model init moves the weight to device.

num_groups

num_groups: int

Number of blocked connections from input channels to output channels.

padding

padding: tuple[int, int, int, int, int, int]

Controls the amount of padding applied before and after the input for depth, height, and width dimensions.

Format: (pad_front, pad_back, pad_top, pad_bottom, pad_left, pad_right).

permute

permute: bool = False

bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, depth, height, width) Max API order: (depth, height, width, in_channels / num_groups, out_channels).

stride

stride: tuple[int, int, int]

Controls the stride for the cross-correlation.

Was this page helpful?