Convolution & Pooling¶

Dimension-agnostic spatial convolutions, poolings, folding, and unfolding.

class taktiny.nn.Conv(in_channels, out_channels, kernel_size, *, stride=1, padding=0, dilation=1, groups=1, pad_mode='zeros', bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: Module

Applies an N-dimensional convolution to channels-last inputs.

The spatial rank is inferred from kernel_size. Scalar channel counts behave like a conventional convolution. Tuple-shaped channels are stored as structured trailing axes and flattened only for the underlying JAX convolution. Thus an input of shape [batch, *spatial, *in_channels] produces [batch, *output_spatial, *out_channels].

When groups is greater than one, groups partition the first input and output channel axes. Both first channel-axis sizes must therefore be divisible by groups.

padding and pad_mode control different aspects of boundary handling. padding determines how many elements are added before and after each spatial axis, while pad_mode determines the values used for those elements. With the default pad_mode='zeros', padding is handled directly by the convolution and the added values are zero. The other modes explicitly extend the input before applying a VALID convolution: 'reflect' mirrors the input without repeating its edge, 'replicate' repeats the edge value, and 'circular' wraps values from the opposite edge.

padding accepts 'VALID' for no automatic padding and 'SAME' or 'SAME_LOWER' for the padding required to produce ceil(input_size / stride) positions along each spatial axis. When the total padding is odd, 'SAME' places the extra element after the input, while 'SAME_LOWER' places it before the input. String padding is supported only with pad_mode='zeros'. Numeric padding can be expressed in several ways:

An integer applies that amount symmetrically to every spatial axis.
A sequence of integers supplies one symmetric amount per spatial axis.
A sequence of (before, after) pairs supplies asymmetric padding for every spatial axis.
For a one-dimensional convolution, a two-integer sequence is interpreted directly as (before, after).

For example, padding=2 pads every axis by two elements on each side; for a two-dimensional convolution, padding=(1, 2) pads the first axis by one and the second by two on each side, while padding=((1, 0), (2, 3)) specifies every side independently.

Parameters:

in_channels (int | Sequence[int]) – Shape of the trailing input-channel axes.
out_channels (int | Sequence[int]) – Shape of the trailing output-channel axes.
kernel_size (int | Sequence[int]) – Size of the spatial convolution window.
stride (int | Sequence[int]) – Step of the convolution window.
padding (str | int | Sequence[int | tuple[int, int]]) – Spatial padding geometry. Use 'VALID' for no automatic padding, 'SAME' or 'SAME_LOWER' to preserve the stride-scaled spatial size, a non-negative integer for symmetric padding on every axis, one integer per axis for per-axis symmetric padding, or a sequence of n (before, after) pairs—one for each spatial axis—for asymmetric padding. Defaults to 0.
dilation (int | Sequence[int]) – Spacing between kernel elements.
groups (int) – Number of feature groups.
pad_mode (str) – How values outside the input boundary are produced. One of 'zeros', 'reflect', 'replicate', or 'circular'. Nonzero modes require explicit numeric padding; 'SAME', 'SAME_LOWER', and 'VALID' can only be used with 'zeros'. Defaults to 'zeros'.
bias (bool) – Whether to add a learnable output bias.
dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to the parameter initializers.
rngs (Rngs) – Random number generator used to initialize parameters.
kernel_initializer (Initializer) – Function used to initialize the kernel.
bias_initializer (Initializer) – Function used to initialize the bias.
quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for the kernel.
dot_general (ConvGeneralDilated | None) – Optional drop-in convolution callable. The name is kept for compatibility with other parameterized modules.
axis_names (tuple[str | None, ...] | None) – Optional logical names for every kernel axis.
partition_spec (P | None) – Optional partition specification for the kernel.
kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the kernel parameter.
bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the bias parameter.
precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Convolution precision forwarded to the convolution callable.
preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred accumulation and result data type.

Examples

Apply a one-dimensional convolution to a channels-last batch while preserving its spatial length:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> conv = nn.Conv(
...     3, 8, kernel_size=3, padding='SAME', rngs=nn.Rngs(0)
... )
>>> x = jnp.ones((4, 16, 3))
>>> conv(x).shape
(4, 16, 8)

Structured channel shapes remain visible in both the input and output. This example applies an unbatched two-dimensional convolution:

>>> conv = nn.Conv(
...     (2, 3),
...     (4, 5),
...     kernel_size=(3, 3),
...     padding='SAME',
...     rngs=nn.Rngs(1),
... )
>>> x = jnp.ones((8, 8, 2, 3))
>>> conv(x).shape
(8, 8, 4, 5)

Nonzero boundary modes require explicit numeric padding. Here the spatial input is reflected by one element on each side:

>>> conv = nn.Conv(
...     1,
...     4,
...     kernel_size=3,
...     padding=1,
...     pad_mode='reflect',
...     rngs=nn.Rngs(2),
... )
>>> x = jnp.ones((6, 1))
>>> conv(x).shape
(6, 4)

class taktiny.nn.ConvTranspose(in_channels, out_channels, kernel_size, *, stride=1, padding=0, dilation=1, groups=1, output_padding=0, bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]¶

Bases: Module

Applies an N-dimensional transposed convolution.

Inputs use the same channels-last layout as Conv. An unbatched input has shape [*spatial, *in_channels] and a batched input has shape [batch, *spatial, *in_channels]. Structured channel axes are flattened only for the underlying convolution and restored in the output.

Numeric padding describes the padding of the corresponding forward convolution. Increasing it crops more values from the transposed output. For each spatial axis, the output size is

(input - 1) * stride - before - after + effective_kernel + output_padding,

where effective_kernel = dilation * (kernel_size - 1) + 1. 'VALID' produces the full transposed-convolution output. 'SAME' and 'SAME_LOWER' produce input_size * stride positions and differ only in which boundary receives an odd extra amount. String padding cannot be combined with output_padding.

When groups is greater than one, groups partition the first input and output channel axes. Both first channel-axis sizes must be divisible by groups.

Parameters:

in_channels (int | Sequence[int]) – Shape of the trailing input-channel axes.
out_channels (int | Sequence[int]) – Shape of the trailing output-channel axes.
kernel_size (int | Sequence[int]) – Size of the spatial convolution window.
stride (int | Sequence[int]) – Factor by which each input position expands the spatial output.
padding (str | int | Sequence[int | tuple[int, int]]) – Forward-convolution padding to remove from the transposed output. Accepts 'VALID', 'SAME', 'SAME_LOWER', a non-negative integer, one symmetric integer per spatial axis, or one (before, after) pair per spatial axis. Defaults to 0.
dilation (int | Sequence[int]) – Spacing between kernel elements.
groups (int) – Number of independent channel groups.
output_padding (int | Sequence[int]) – Additional size added to the end of each output spatial axis. It resolves shape ambiguity when stride > 1 and does not pad the output with values. Each amount must be smaller than either its stride or dilation. Defaults to 0.
bias (bool) – Whether to add a learnable output bias.
dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to the parameter initializers.
rngs (Rngs) – Random number generator used to initialize parameters.
kernel_initializer (Initializer) – Function used to initialize the kernel.
bias_initializer (Initializer) – Function used to initialize the bias.
quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for the kernel.
dot_general (ConvGeneralDilated | None) – Optional replacement for conv_general_dilated.
axis_names (tuple[str | None, ...] | None) – Optional logical names for every kernel axis.
partition_spec (P | None) – Optional partition specification for the kernel.
kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the kernel parameter.
bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the bias parameter.
precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Convolution precision forwarded to the convolution callable.
preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred accumulation and result data type.

Variables:

kernel – Learnable kernel with shape (*kernel_size, *in_channels, *grouped_out_channels).
bias – Learnable bias with shape out_channels, or None.

Examples

Upsample a one-dimensional input by a factor of two:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> conv = nn.ConvTranspose(
...     3, 4, kernel_size=3, stride=2, rngs=nn.Rngs(0)
... )
>>> conv(jnp.ones((5, 3))).shape
(11, 4)

Structured channel shapes are preserved:

>>> conv = nn.ConvTranspose(
...     (2, 3),
...     (4, 5),
...     kernel_size=(2, 2),
...     stride=2,
...     rngs=nn.Rngs(1),
... )
>>> conv(jnp.ones((3, 3, 2, 3))).shape
(6, 6, 4, 5)

class taktiny.nn.MaxPool(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)[source]¶

Bases: _SpatialOp

Take maxima over channels-last spatial windows.

Inputs are (*spatial, channels) or (batch, *spatial, channels), with one trailing channel axis. Scalars specify 1-D spatial shapes; sequences specify the spatial rank. Spatial and channel dimensions must be nonempty. __call__ accepts out_sharding for the final output.

Parameters:

kernel_size (int | Sequence[int]) – Positive spatial window sizes.
stride (int | Sequence[int] | None) – Positive strides; None uses kernel_size.
padding (str | int | Sequence[int | tuple[int, int]]) – VALID, SAME, SAME_LOWER or nonnegative symmetric/asymmetric numeric padding, as in Conv. Padding uses the dtype’s minimum identity (negative infinity for floating-point inputs), not zero.
dilation (int | Sequence[int]) – Positive spacing within each pooling window; defaults to 1.
return_indices (bool) – Return (values, indices) when True.
ceil_mode (bool) – Include a final partial window by extending right padding.

Indices are int32 row-major spatial offsets, excluding batch and channel. Ties choose the smallest spatial offset. NaNs propagate; tied NaNs choose the first offset. Fully padded windows use the maximum int32 index sentinel. Complex inputs are unsupported. With indices enabled, out_sharding is applied to both result arrays.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> values, indices = nn.MaxPool(2, return_indices=True)(jnp.arange(4.)[:, None])
>>> values[:, 0].tolist(), indices[:, 0].tolist()
([1.0, 3.0], [1, 3])

class taktiny.nn.AvgPool(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None)[source]¶

Bases: _SpatialOp

Average channels-last spatial windows.

Parameters:

kernel_size (int | Sequence[int]) – Positive window sizes.
stride (int | Sequence[int] | None) – Positive window strides; None uses kernel_size.
padding (str | int | Sequence[int | tuple[int, int]]) – VALID, SAME, SAME_LOWER or explicit nonnegative padding. Outside-input values contribute zero to each window sum.
ceil_mode (bool) – Permit a final partial window with extra right padding.
count_include_pad (bool) – Include configured zero padding in the divisor. Extra padding introduced solely by ceil_mode is never counted.
divisor_override (int | None) – Optional positive integer divisor for every window.

Integer and boolean inputs are promoted to float32. No RNGs or parameters are used. With count_include_pad=False, the divisor counts actual input elements; configurations containing wholly padded windows have zero count.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.AvgPool(2)(jnp.arange(4.)[:, None])[:, 0].tolist()
[0.5, 2.5]

class taktiny.nn.AdaptiveMaxPool(output_size, return_indices=False)[source]¶

Bases: _SpatialOp

Pool variable-size spatial bins to a requested output shape.

Parameters:

output_size (int | Sequence[int | None]) – Positive target sizes. A None entry preserves that input spatial dimension. A scalar specifies one spatial dimension.
return_indices (bool) – Return values and int32 row-major spatial offsets.

Bin i spans floor(i * input / output) through ceil((i+1) * input / output), excluding the end. Bins may overlap; output sizes may exceed input sizes. Each channel is pooled separately. Ties choose the first element in the bin. Complex inputs are unsupported. out_sharding applies to both arrays when return_indices=True.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.AdaptiveMaxPool(3)(jnp.arange(6.)[:, None])[:, 0].tolist()
[1.0, 3.0, 5.0]

class taktiny.nn.AdaptiveAvgPool(output_size)[source]¶

Bases: _SpatialOp

Average variable-size spatial bins to a requested output shape.

Parameters:: output_size (int | Sequence[int | None]) – Positive target sizes; None entries preserve the matching input dimensions. Scalars specify one spatial dimension.

Bin i spans floor(i * input / output) through ceil((i+1) * input / output), excluding the end. Bins may overlap and each is divided by its own number of elements. Output sizes may exceed input sizes. Integer and boolean inputs promote to float32. This module has no parameters or randomness.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.AdaptiveAvgPool((2, None))(jnp.ones((4, 3, 1))).shape
(2, 3, 1)

class taktiny.nn.FractionalMaxPool(kernel_size, output_size=None, output_ratio=None, return_indices=False, random_samples=None, *, rngs=None)[source]¶

Bases: _SpatialOp

Max-pool windows placed on a fractional, optionally random spatial grid.

Parameters:

kernel_size (int | Sequence[int]) – Positive window sizes.
output_size (int | Sequence[int] | None) – Positive target spatial sizes; mutually exclusive with output_ratio. Exactly one must be provided.
output_ratio (float | Sequence[float] | None) – Finite ratios in (0, 1], scalar or per-axis. Target sizes are max(1, floor(input_size * ratio)). Windows must fit the input.
return_indices (bool) – Also return int32 flattened spatial maximum indices.
random_samples (Array | Sequence[float] | None) – Optional fixed array/sequence of shape (spatial_rank,) with values in [0, 1). Eager values are validated; traced values must satisfy this domain. One grid is shared by batches/channels.
rngs (Rngs | None) – Explicit runtime stream. If fixed samples are absent, consumes a fresh key on each call; None uses get_context_rng at call time. A missing stream raises an error. Fixed samples consume no RNG.

Sampling also occurs in eval mode: this is grid sampling, not dropout. To retain the former deterministic default grid, pass random_samples=(0.5,) for 1-D, or one 0.5 per spatial axis. With owned RNGs under jit, pass the module in and return its updated state. Values and optional indices both receive out_sharding. Complex inputs are unsupported.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> pool = nn.FractionalMaxPool(2, output_size=3, random_samples=(0.5,))
>>> pool(jnp.arange(6.)[:, None])[:, 0].tolist()
[1.0, 3.0, 5.0]

class taktiny.nn.LPPool(norm_type, kernel_size, stride=None, ceil_mode=False)[source]¶

Bases: _SpatialOp

Compute (sum(abs(x)**p))**(1/p) over spatial windows.

Parameters:

norm_type (float) – Finite positive exponent p. This is a true norm when p>=1; values below one are allowed as a power aggregation.
kernel_size (int | Sequence[int]) – Positive spatial window sizes.
stride (int | Sequence[int] | None) – Positive strides; None uses kernel_size.
ceil_mode (bool) – Permit a final partial window, treating missing values as zero.

This is a sum-based Lp aggregation, not a power mean: there is no division by window volume. Integer inputs promote to float32; complex inputs use real magnitudes. The zero-total output uses a zero derivative convention.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.LPPool(2, 2)(jnp.array([[3.], [4.]]))[:, 0].tolist()
[5.0]

class taktiny.nn.MaxUnpool(kernel_size, stride=None, padding=0, dilation=1)[source]¶

Bases: _SpatialOp

Scatter pooled values back to their flattened spatial indices.

Parameters:

kernel_size (int | Sequence[int]) – Original positive pooling window sizes.
stride (int | Sequence[int] | None) – Original strides; None uses kernel_size.
padding (int | Sequence[int | tuple[int, int]]) – Original explicit numeric padding; strings are unsupported.
dilation (int | Sequence[int]) – Original positive window dilation; defaults to 1.

__call__(x, indices, output_size=None, out_sharding=None) requires integer indices with the same shape as x. output_size contains spatial dimensions only and overrides the shape inferred from kernel, stride, dilation and padding. Supply it when ceil-mode or strided pooling made the original size ambiguous. Out-of-range indices are discarded, not wrapped. Unfilled positions are zero. Repeated indices have unspecified write order, so this is only a partial inverse of MaxPool, not a reconstruction of lost values.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> x = jnp.arange(4.)[:, None]
>>> values, indices = nn.MaxPool(2, return_indices=True)(x)
>>> nn.MaxUnpool(2)(values, indices)[:, 0].tolist()
[0.0, 1.0, 0.0, 3.0]

class taktiny.nn.Padding(padding, mode='constant', value=0.0)[source]¶

Bases: _SpatialOp

Pad arbitrary array axes using jax.numpy.pad conventions.

Unlike the pooling modules, every axis is eligible, including batch and channels. No channels-last interpretation or batch insertion is performed.

Parameters:

padding (int | Sequence[int] | Sequence[tuple[int, int]]) – Nonnegative integer applied to both ends of every axis; a flat (before, after) pair broadcast to every axis; or one such pair per array axis. A one-element sequence broadcasts symmetrically.
mode (str) – constant, edge, reflect, symmetric or wrap. Aliases zeros, replicate and circular map to constant, edge and wrap.
value (float) – Fill value for constant mode only; defaults to 0.

reflect excludes the edge element when mirroring; symmetric repeats it. Numeric padding is normalized to immutable tuples at construction. __call__ accepts out_sharding for the final padded array.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Padding((1, 2), value=9)(jnp.array([1, 2])).tolist()
[9, 1, 2, 9, 9]

class taktiny.nn.Fold(output_size, kernel_size, dilation=1, padding=0, stride=1)[source]¶

Bases: _SpatialOp

Overlap-add flattened sliding patches into a channels-last array.

Parameters:

output_size (int | Sequence[int]) – Positive target spatial dimensions; determines spatial rank.
kernel_size (int | Sequence[int]) – Positive window sizes, scalar or per-axis.
dilation (int | Sequence[int]) – Positive spacing within windows; defaults to 1.
padding (str | int | Sequence[int | tuple[int, int]]) – The same VALID/SAME/SAME_LOWER or numeric padding used by Unfold.
stride (int | Sequence[int]) – Positive window strides; defaults to 1.

Input is (windows, patch_width) or (batch, windows, patch_width). Patch layout must match Unfold: channel followed by flattened kernel dimensions. Output is (*output_size, channels), optionally with a batch axis. Overlaps are summed and padded positions are discarded, including negative coordinates; they never wrap around the output. To invert Unfold, divide by the overlap counts where those counts are nonzero. __call__ accepts out_sharding. Patch width must be a positive multiple of kernel volume.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Fold(4, 2)(jnp.ones((3, 2)))[:, 0].tolist()
[1.0, 2.0, 2.0, 1.0]

class taktiny.nn.Unfold(kernel_size, dilation=1, padding=0, stride=1)[source]¶

Bases: _SpatialOp

Extract channels-last sliding windows into flattened patches.

Parameters:

kernel_size (int | Sequence[int]) – Positive spatial window sizes (GenericShape).
dilation (int | Sequence[int]) – Positive spacing within each window, scalar or per-axis.
padding (str | int | Sequence[int | tuple[int, int]]) – VALID, SAME, SAME_LOWER, a nonnegative symmetric integer, per-axis integers, or per-axis (before, after) pairs. For 1-D, a flat pair denotes asymmetric padding. Padding contributes zeros.
stride (int | Sequence[int]) – Positive window strides; defaults to 1.

Returns (windows, patch_width) or (batch, windows, patch_width). Windows are flattened in spatial row-major order. Within each patch, channel comes before the kernel axes: patch_width = channels * prod(kernel_size). Raises ValueError when effective kernel and padding produce no windows. Fold sums overlapping patches; it is not automatically an inverse.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Unfold(2)(jnp.ones((4, 1))).shape
(3, 2)