Spatial

spatial

PatchEmbed2d

PatchEmbed2d(in_channels: int = 3, embed_dim: int = 768, patch_size: int | tuple = 16, bias: bool = True)

Bases: Module

2D Patch Embedding using Conv2d.

Projects image patches into an embedding space.

Parameters:

Name	Type	Description	Default
in_channels	int	Input image channels.	3
embed_dim	int	Output embedding dimension.	768
patch_size	int | tuple	Patch size (int or (h, w)).	16
bias	bool	Whether to use bias in the projection.	True

PatchEmbed3d

PatchEmbed3d(in_channels: int = 3, embed_dim: int = 768, patch_size: int | tuple = (2, 16, 16), bias: bool = True)

Bases: Module

3D Patch Embedding using Conv3d.

Projects video patches into an embedding space.

Parameters:

Name	Type	Description	Default
in_channels	int	Input channels.	3
embed_dim	int	Output embedding dimension.	768
patch_size	int | tuple	Patch size (int or (d, h, w)).	(2, 16, 16)
bias	bool	Whether to use bias in the projection.	True

avg_pool1d

avg_pool1d(x: array, kernel_size: int, stride: int | None = None) -> array

1D average pooling.

Equivalent to torch.nn.functional.avg_pool1d.

Parameters:

Name	Type	Description	Default
x	array	(B, L, C) input tensor.	required
kernel_size	int	Pooling window size.	required
stride	int | None	Pooling stride. Defaults to kernel_size.	None

Returns:

Type	Description
array	(B, L_out, C) pooled tensor where L_out = (L - kernel_size) // stride + 1.

interpolate_3d

interpolate_3d(x: array, size: tuple) -> array

Nearest-neighbor interpolation for 5D (B,D,H,W,C) tensors.

Convenience wrapper around interpolate_nearest for video tensors.

Parameters:

Name	Type	Description	Default
x	array	(B, D, H, W, C) input tensor.	required
size	tuple	Target (D, H, W).	required

Returns:

Type	Description
array	(B, target_D, target_H, target_W, C) interpolated tensor.

interpolate_nearest

interpolate_nearest(x: array, size: int | tuple[int, ...] | None = None, scale_factor: float | None = None) -> array

Nearest-neighbor interpolation for N-dimensional spatial tensors.

Supports 3D (B,L,C), 4D (B,H,W,C), and 5D (B,D,H,W,C) inputs. Operates on all spatial dims (all except batch and channel).

Parameters:

Name	Type	Description	Default
x	array	Input tensor in channels-last format.	required
size	int | tuple[int, ...] | None	Target spatial size. Mutually exclusive with scale_factor.	None
scale_factor	float | None	Scale factor for all spatial dims.	None

Returns:

Type	Description
array	Interpolated tensor.

replicate_pad

replicate_pad(x: array, pad_widths: list) -> array

Pad a tensor by replicating edge values.

Equivalent to torch.nn.functional.pad(x, ..., mode="replicate").

Parameters:

Name	Type	Description	Default
x	array	Input tensor of any shape.	required
pad_widths	list	List of (before, after) padding per dimension. Length must match x.ndim.	required

Returns:

Type	Description
array	Padded tensor.

patchify

patchify(x: array, patch_size: tuple | int) -> array

Convert spatial input into a sequence of flattened patches.

Parameters:

Name	Type	Description	Default
x	array	Input tensor. 2D: (B, H, W, C) -> patches of (ph, pw) 3D: (B, D, H, W, C) -> patches of (pd, ph, pw)	required
patch_size	tuple | int	Patch dimensions. Int for uniform, tuple for per-dim.	required

Returns:

Type	Description
array	(B, num_patches, patch_dim) where patch_dim = prod(patch_size) * C.

unpatchify

unpatchify(x: array, patch_size: tuple | int, shape: tuple) -> array

Reconstruct spatial tensor from a sequence of patches.

Parameters:

Name	Type	Description	Default
x	array	(B, num_patches, patch_dim) patch sequence.	required
patch_size	tuple | int	Patch dimensions.	required
shape	tuple	Target spatial shape without batch and channels. 2D: (H, W), 3D: (D, H, W).	required

Returns:

Type	Description
array	Reconstructed tensor: (B, *shape, C).

pixel_unshuffle

pixel_unshuffle(x: array, downscale_factor: int) -> array

Rearrange spatial dimensions into channels (inverse of pixel_shuffle).

Channels-last equivalent of torch.nn.functional.pixel_unshuffle using the PT channel-ordering convention (C, patch_row, patch_col). Round-trip with :func:pixel_shuffle (same upscale/downscale factor) is the identity.

Parameters:

Name	Type	Description	Default
x	array	(B, H, W, C) tensor. H and W must be divisible by downscale_factor.	required
downscale_factor	int	Factor by which to decrease spatial resolution.	required

Returns:

Type	Description
array	(B, H/r, W/r, C*r^2) where r = downscale_factor.

upsample_bilinear

upsample_bilinear(x: array, scale_factor: int = 2) -> array

Bilinear upsampling for 2D spatial tensors (B, H, W, C).

Uses the formula: output[i,j] = weighted average of 4 nearest input pixels.

Parameters:

Name	Type	Description	Default
x	array	(B, H, W, C) input tensor.	required
scale_factor	int	Integer upsampling factor.	2

Returns:

Type	Description
array	(B, Hscale_factor, Wscale_factor, C) upsampled tensor.

upsample_nearest

upsample_nearest(x: array, scale_factor: int = 2) -> array

Nearest-neighbor upsampling for spatial tensors.

Parameters:

Name	Type	Description	Default
x	array	Input tensor (B, H, W, C) or (B, D, H, W, C).	required
scale_factor	int	Integer upsampling factor.	2

Returns:

Type	Description
array	Upsampled tensor.

interpolate

Interpolation operations missing from MLX core.

Equivalent to torch.nn.functional.interpolate for various modes.

interpolate_nearest

interpolate_nearest(x: array, size: int | tuple[int, ...] | None = None, scale_factor: float | None = None) -> array

Nearest-neighbor interpolation for N-dimensional spatial tensors.

Supports 3D (B,L,C), 4D (B,H,W,C), and 5D (B,D,H,W,C) inputs. Operates on all spatial dims (all except batch and channel).

Parameters:

Name	Type	Description	Default
x	array	Input tensor in channels-last format.	required
size	int | tuple[int, ...] | None	Target spatial size. Mutually exclusive with scale_factor.	None
scale_factor	float | None	Scale factor for all spatial dims.	None

Returns:

Type	Description
array	Interpolated tensor.

interpolate_3d

interpolate_3d(x: array, size: tuple) -> array

Nearest-neighbor interpolation for 5D (B,D,H,W,C) tensors.

Convenience wrapper around interpolate_nearest for video tensors.

Parameters:

Name	Type	Description	Default
x	array	(B, D, H, W, C) input tensor.	required
size	tuple	Target (D, H, W).	required

Returns:

Type	Description
array	(B, target_D, target_H, target_W, C) interpolated tensor.

avg_pool1d

avg_pool1d(x: array, kernel_size: int, stride: int | None = None) -> array

1D average pooling.

Equivalent to torch.nn.functional.avg_pool1d.

Parameters:

Name	Type	Description	Default
x	array	(B, L, C) input tensor.	required
kernel_size	int	Pooling window size.	required
stride	int | None	Pooling stride. Defaults to kernel_size.	None

Returns:

Type	Description
array	(B, L_out, C) pooled tensor where L_out = (L - kernel_size) // stride + 1.

pad

Padding operations missing from MLX core.

MLX's mx.pad only supports constant and reflect modes. This module adds replicate padding (edge padding).

replicate_pad

replicate_pad(x: array, pad_widths: list) -> array

Pad a tensor by replicating edge values.

Equivalent to torch.nn.functional.pad(x, ..., mode="replicate").

Parameters:

Name	Type	Description	Default
x	array	Input tensor of any shape.	required
pad_widths	list	List of (before, after) padding per dimension. Length must match x.ndim.	required

Returns:

Type	Description
array	Padded tensor.

patch

Patchify/Unpatchify operations and patch embedding layers.

PatchEmbed2d

PatchEmbed2d(in_channels: int = 3, embed_dim: int = 768, patch_size: int | tuple = 16, bias: bool = True)

Bases: Module

2D Patch Embedding using Conv2d.

Projects image patches into an embedding space.

Parameters:

Name	Type	Description	Default
in_channels	int	Input image channels.	3
embed_dim	int	Output embedding dimension.	768
patch_size	int | tuple	Patch size (int or (h, w)).	16
bias	bool	Whether to use bias in the projection.	True

PatchEmbed3d

PatchEmbed3d(in_channels: int = 3, embed_dim: int = 768, patch_size: int | tuple = (2, 16, 16), bias: bool = True)

Bases: Module

3D Patch Embedding using Conv3d.

Projects video patches into an embedding space.

Parameters:

Name	Type	Description	Default
in_channels	int	Input channels.	3
embed_dim	int	Output embedding dimension.	768
patch_size	int | tuple	Patch size (int or (d, h, w)).	(2, 16, 16)
bias	bool	Whether to use bias in the projection.	True

patchify

patchify(x: array, patch_size: tuple | int) -> array

Convert spatial input into a sequence of flattened patches.

Parameters:

Name	Type	Description	Default
x	array	Input tensor. 2D: (B, H, W, C) -> patches of (ph, pw) 3D: (B, D, H, W, C) -> patches of (pd, ph, pw)	required
patch_size	tuple | int	Patch dimensions. Int for uniform, tuple for per-dim.	required

Returns:

Type	Description
array	(B, num_patches, patch_dim) where patch_dim = prod(patch_size) * C.

unpatchify

unpatchify(x: array, patch_size: tuple | int, shape: tuple) -> array

Reconstruct spatial tensor from a sequence of patches.

Parameters:

Name	Type	Description	Default
x	array	(B, num_patches, patch_dim) patch sequence.	required
patch_size	tuple | int	Patch dimensions.	required
shape	tuple	Target spatial shape without batch and channels. 2D: (H, W), 3D: (D, H, W).	required

Returns:

Type	Description
array	Reconstructed tensor: (B, *shape, C).

pixel_shuffle

Pixel Shuffle / Unshuffle operations for sub-pixel convolution.

pixel_shuffle

pixel_shuffle(x: array, upscale_factor: int) -> array

Rearrange channels into spatial dimensions (sub-pixel convolution).

Channels-last equivalent of torch.nn.functional.pixel_shuffle. PyTorch flattens the channel axis as (out_channels, patch_row, patch_col) so channel j of a C = oc * r * r input maps to (c=j // r**2, pr=(j // r) % r, pc=j % r). Any layout that reverses the (oc, r, r) ordering (e.g. (r, r, oc)) shuffles output channels and silently produces checkerboard artefacts downstream — see the mlx-porting skill's common-pitfalls #7 for a past burn case.

Parameters:

Name	Type	Description	Default
x	array	(B, H, W, C) tensor where C must be divisible by upscale_factor^2.	required
upscale_factor	int	Factor by which to increase spatial resolution.	required

Returns:

Type	Description
array	(B, Hr, Wr, C/r^2) where r = upscale_factor.

pixel_unshuffle

pixel_unshuffle(x: array, downscale_factor: int) -> array

Rearrange spatial dimensions into channels (inverse of pixel_shuffle).

Channels-last equivalent of torch.nn.functional.pixel_unshuffle using the PT channel-ordering convention (C, patch_row, patch_col). Round-trip with :func:pixel_shuffle (same upscale/downscale factor) is the identity.

Parameters:

Name	Type	Description	Default
x	array	(B, H, W, C) tensor. H and W must be divisible by downscale_factor.	required
downscale_factor	int	Factor by which to decrease spatial resolution.	required

Returns:

Type	Description
array	(B, H/r, W/r, C*r^2) where r = downscale_factor.

upsample

Upsampling operations.

upsample_nearest

upsample_nearest(x: array, scale_factor: int = 2) -> array

Nearest-neighbor upsampling for spatial tensors.

Parameters:

Name	Type	Description	Default
x	array	Input tensor (B, H, W, C) or (B, D, H, W, C).	required
scale_factor	int	Integer upsampling factor.	2

Returns:

Type	Description
array	Upsampled tensor.

upsample_bilinear

upsample_bilinear(x: array, scale_factor: int = 2) -> array

Bilinear upsampling for 2D spatial tensors (B, H, W, C).

Uses the formula: output[i,j] = weighted average of 4 nearest input pixels.

Parameters:

Name	Type	Description	Default
x	array	(B, H, W, C) input tensor.	required
scale_factor	int	Integer upsampling factor.	2

Returns:

Type	Description
array	(B, Hscale_factor, Wscale_factor, C) upsampled tensor.