Pooling operations
Pooling operations reduce sparse feature rows through a kernel relation or
through batch metadata.
Local pooling computes:
\[y_{o,c} = \operatorname{reduce}_{e:o_e=o} x_{i_e,c}.\]
sum and avg accept empty output rows as zero-valued reductions. max
requires at least one contributing row for every output row. Global pooling
uses batch_counts metadata from the input sparse tensor and returns a dense
(B, C) MLX array.
pool_transpose3d reverses sparse resolution by averaging each coarse row
over a transposed kernel relation. Without a target it generates fine support;
with a target tensor it preserves that tensor’s coordinates and emits zeros for
target rows with no contributors.
trilinear_upsample3d performs normalized separable interpolation instead of
uniform averaging. Its optional target follows the same support contract as
pooling transpose.
Related pages
-
mlx_lattice.ops.pool.avg_pool3d(x, *, kernel_size=2, stride=2, padding=0, dilation=1)[source]
Apply local sparse average pooling.
The result feature at each output row is the sparse sum divided by the
number of contributing relation edges for that output row.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.pool.global_avg_pool(x, *, batch_size=None)[source]
Average features independently for each batch.
Empty batches produce zero rows. batch_counts is used when present;
otherwise batches are reduced from the coordinate batch column. Pass
batch_size to preserve trailing empty batches.
- Return type:
array
- Parameters:
-
-
mlx_lattice.ops.pool.global_max_pool(x, *, batch_size=None)[source]
Max-reduce features independently for each batch.
Empty batches are rejected because max has no neutral finite sparse row.
batch_counts is used when present; otherwise batches are reduced from
the coordinate batch column. Pass batch_size to validate explicit empty
batches.
- Return type:
array
- Parameters:
-
-
mlx_lattice.ops.pool.global_pool(x, *, mode='sum', batch_size=None)[source]
Reduce sparse features independently for each batch.
mode selects sum, avg, or max. batch_size=None infers
the number of dense output rows from sparse batch metadata or coordinate
values. Explicit batch_size preserves trailing empty batches for
sum and avg and validates them for max.
- Return type:
array
- Parameters:
-
-
mlx_lattice.ops.pool.global_sum_pool(x, *, batch_size=None)[source]
Sum features independently for each batch.
Returns a dense (B, C) MLX array. batch_counts is used when
present; otherwise batches are reduced from the coordinate batch column.
Pass batch_size to preserve trailing empty batches.
- Return type:
array
- Parameters:
-
-
mlx_lattice.ops.pool.max_pool3d(x, *, kernel_size=2, stride=2, padding=0, dilation=1)[source]
Apply local sparse max pooling.
The result feature at each output row is the channel-wise maximum over
contributing input rows in the sparse kernel relation.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.pool.pool3d(x, *, mode='sum', kernel_size=2, stride=2, padding=0, dilation=1)[source]
Apply local sparse 3D pooling with sum, max, or avg mode.
Local pooling builds a forward kernel relation and reduces input features
that contribute to each output coordinate. The output sparse stride is
x.stride * stride. Native local pooling accepts float32 and,
on Metal, inference-only float16 features. Metal routes additionally
require int32 coordinates.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.pool.pool_transpose3d(x, target=None, *, kernel_size=2, stride=2, padding=0, dilation=1)[source]
Average coarse features onto generated or explicit fine support.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.pool.sum_pool3d(x, *, kernel_size=2, stride=2, padding=0, dilation=1)[source]
Apply local sparse sum pooling.
The result feature at each output row is the sum of all contributing input
rows in the sparse kernel relation.
- Return type:
SparseTensor
- Parameters:
-
-
mlx_lattice.ops.pool.trilinear_upsample3d(x, target=None, *, stride=2)[source]
Upsample sparse features with normalized trilinear interpolation.
- Return type:
SparseTensor
- Parameters:
-