Skip to content

Resampling

Resampler¤

Resamples inputs using different interpolation methods. Useful for architecures like the UNO, where data needs to be explicitly up- and downsampled as it flows through the network.

neojax.nn.resample.Resampler ¤

Resample functions via n-dimensional interpolation.

Uses linear interpolation for 1D inputs and bicubic interpolation for 2D inputs.

For 3D and higher-dimensional inputs, spectral interpolation is used as follows:

  1. Input is transformed to frequency domain using real n-dim FFT.
  2. Frequency components are resized.
  3. Inverse FFT to spatial domain, resulting in smooth, alias-free interpolation.

Parameters:

Name Type Description Default
input_shape Sequence[int]

Shape of input including channel dimension, namely (c, d1, ..., dN).

 required
res_scale float | int | Sequence[int | float]

Scaling factor along each of the dimensions in axes. If res_scale is scalar, then isotropic scaling is performed.

 required
axes int | tuple[int, ...] | None

Optional axes/spatial dimension along which interpolation will be performed. Channel dimension is treated as 0-th dimension.

 required
output_shape Sequence[int] | None

Optional output shape. Default is None. If output_shape is passed, res_scale is ignored.

 None
Internal Attributes
  • res_scale (tuple[int, ...]): Scaling factor per spatial dimension.
  • axes (tuple[int, ...]): Axes along which to interpolate.
  • output_shape (tuple[int, ...]): Output shape including channel dimension.
  • input_shape (tuple[int, ...]): Shape of input including channel dimension.
  • output_modes (tuple[int, ...]): Output modes to resize frequency components with.

Examples:

import jax.random as jr
from neojax.nn import Resampler

key = jr.key(0)
x = jr.normal(key=key, shape=(2, 10))
resampler_1d = Resampler(input_shape=x.shape, res_scale=2, axes=1)

# out.shape = (2, 20)
out = resampler_1d(x)

Info

This class provides the same functionality as the resample function in the neuraloperator library. It imitates the align_corners behavior using jax.image.scale_and_translate for 2d inputs. See here for details.

__call__(x: Float[Array, 'c ...']) -> Float[Array, 'c ...'] ¤

Resamples n-dimensional input array.

Parameters:

Name Type Description Default
x Float[Array, 'c ...']

Inputs shaped (c, d1, ..., dN).

 required

Returns:

Type Description
Float[Array, 'c ...']

Resampled outputs.