Resampling

Spatial resampling, downsampling, and upsampling layers.

class taktiny.nn.Downsample(size=None, scale_factor=None, method='linear', antialias=True, *, precision=Precision.HIGHEST)[source]

Bases: _Resize

Reduce channels-last spatial dimensions by a divisive factor.

For n spatial dimensions, accepts (*spatial, channels) or (batch, *spatial, channels). Scalar size/scale_factor describes 1-D input; use a sequence such as (2, 2) for 2-D resizing. Batch and the single trailing channel axis are preserved; N-D feature blocks are not inferred.

Parameters:

  • size (int | Sequence[int | None] | None) – Positive target sizes no larger than the input. None entries preserve individual dimensions. Mutually exclusive with scale_factor.

  • scale_factor (float | Sequence[float] | None) – Finite divisors greater than or equal to one. Each output size is max(1, floor(input_size / factor)). Defaults to 2 for 1-D input when neither size nor scale_factor is supplied.

  • method (str | ResizeMethod) – JAX resize method string or ResizeMethod enum. Defaults to linear; supports nearest, cubic, Lanczos and JAX’s method aliases.

  • antialias (bool) – Filter when reducing spatial dimensions to limit aliasing. Defaults to True. Ignored for nearest-neighbor interpolation.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Interpolation contraction precision; defaults to HIGHEST, matching jax.image.resize. Ignored for nearest interpolation.

__call__ accepts out_sharding for the final output. Nearest preserves the input dtype; other methods follow JAX’s floating-point promotion. Input spatial dimensions must be nonempty. This operation has no RNG or parameters.

Examples

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Downsample(scale_factor=(2, 3))(jnp.ones((2, 8, 9, 4))).shape
(2, 4, 3, 4)
>>> nn.Downsample(size=(3, None))(jnp.ones((7, 5, 4))).shape
(3, 5, 4)
class taktiny.nn.Upsample(size=None, scale_factor=None, method='nearest', antialias=True, *, precision=Precision.HIGHEST)[source]

Bases: _Resize

Resize channels-last spatial dimensions by a multiplicative factor.

For n spatial dimensions, accepts (*spatial, channels) or (batch, *spatial, channels). Scalar size/scale_factor means n=1; a sequence specifies n explicitly. For 2-D images, use e.g. scale_factor=(2, 2). Batch and channel dimensions are preserved. Only one channel axis is supported; trailing N-D feature blocks are not inferred.

Parameters:

  • size (int | Sequence[int | None] | None) – Positive target sizes. None entries in a sequence preserve the corresponding input dimensions. Mutually exclusive with scale_factor.

  • scale_factor (float | Sequence[float] | None) – Finite positive multipliers. Each output size is max(1, floor(input_size * factor)). Defaults to 2 for 1-D input when neither size nor scale_factor is given. Factors below one and smaller explicit sizes remain supported for compatibility.

  • method (str | ResizeMethod) – JAX resize method string or ResizeMethod enum. Defaults to nearest; supports linear, cubic, Lanczos and JAX’s method aliases.

  • antialias (bool) – Apply filtering when shrinking. No effect when enlarging or using nearest-neighbor interpolation. Defaults to True.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Interpolation contraction precision; defaults to HIGHEST, matching jax.image.resize. Ignored for nearest interpolation.

__call__ accepts out_sharding for the final output. Nearest preserves the input dtype; other methods follow JAX’s floating-point promotion. Input spatial dimensions must be nonempty. No parameters or RNGs are used.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> layer = nn.Upsample(scale_factor=(2, 3), method='linear')
>>> layer(jnp.ones((4, 5, 3))).shape
(8, 15, 3)