Filtered Leaky ReLU

Filtered Leaky ReLU

neojax.nn.filtered_lrelu.FilteredLReLU

Filtered leaky ReLU block (n-dimensional).

Computes the alias-free activation layer fused with the preceding/subsequent resampling operator of the CNO (does not include the preceding convolution layer).

Which activation layer and resampling operator to compute can be controlled via block_type:

• "invariant" (I-block/R-block): \(u = 2, d = 2\). s -> 2s -> lReLU -> s
• "upsample" (U-block): \(u = 4, d = 2\). s -> 4s -> lReLU -> 2s
• "downsample" (D-block): \(u = 2, d = 4\). s -> 2s -> lReLU -> s/2

Implements the CUDA logic as in the reference (classic CNO variant), i.e. as in StyleGAN3 in JAX. The filter taps are compile-time constants.

Parameters:

Name	Type	Description	Default
ndim	int	Number of spatial dimensions.	required
channels	int	Number of channels of the input, needed for the bias.	required
h_up	KaiserFIRFilter	Kaiser windowed FIR filter used for upsampling. It needs up * taps taps, with up being the upsampling factor of the block_type, and taps any integer.	required
h_down	KaiserFIRFilter	Kaiser windowed FIR filter used for downsampling. It needs down * taps taps, with down being the downsampling factor of the block_type, and taps any integer.	required
block_type	Literal['invariant', 'upsample', 'downsample']	The type of activation layer + resampling operator block to compute. Must be one of "invariant", "upsample" or "downsample". See class docstring for details.	required
use_bias	bool	Whether to add a learnable channel-wise bias (initialized to zero) before the activation. Default is True.	True
pad_mode	Literal['wrap', 'reflect', 'constant']	Padding mode for the up- and downsampling padding. Use "wrap" for a periodic domain, otherwise "reflect" or "constant" (zero padding as in the reference). Default is "wrap".	'wrap'
negative_slope	float	Negative slope for the leaky ReLU (multiplication factor for x < 0). Default is 0.2.	0.2
gain	float	Factor multiplying the output of the leaky ReLU. Default is \(\sqrt{2}\), as in the reference.	1.4142135623730951

Raises:

Type	Description
ValueError	If block_type or pad_mode are invalid.
ValueError	If negative_slope is less than or equal to zero.
ValueError	If the number of taps of h_up (h_down) is not a multiple of the up- (down-)sampling factor, or if h_down has an odd number of taps.

Internal Attributes

These fields store the internal layer states and weights.

• bias (Float[Array, "c"] | None): The learnable channel-wise bias.
• h_up (KaiserFIRFilter): The upsampling filter.
• h_down (KaiserFIRFilter): The downsampling filter.
• resampling_factors (dict[str, tuple[int, int]]): Dictionary mapping block_type to up- and downsampling factors.
• block_type (str): The type of block.
• pad_mode (str): The padding mode.
• spatial_axes (tuple[int, ...]): The (negative) spatial axes.
• negative_slope (float): The slope of the leaky ReLU.
• gain (float): The gain after the leaky ReLU.

Examples:

import jax.numpy as jnp
from neojax.nn import FilteredLReLU, KaiserFIRFilter

# Filters of an I-block, designed on a unit grid
h_up = KaiserFIRFilter(1 / 2.0001, 6, 0.8 - 1 / 2.0001, 2, 1)
h_down = KaiserFIRFilter(1 / 2.0001, 6, 0.8 - 1 / 2.0001, 2, 1)
act = FilteredLReLU(
    ndim=2,
    channels=3,
    h_up=h_up,
    h_down=h_down,
    block_type="invariant",
)
x = jnp.ones((3, 32, 32))
# out.shape = (3, 32, 32)
out = act(x)

down property

Downsampling factor of the activation rate w.r.t. the output.

up property

Upsampling factor of the activation rate w.r.t the input.

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

Forward pass through the fused sampling and LReLU block.

Parameters:

Name	Type	Description	Default
x	Inexact[Array, 'c ...']	Input signal shaped (c, d1, ..., dN).	required

Returns:

Type	Description
Inexact[Array, 'c ...']	The up- or downsampled signal with applied LReLU activation, shaped (c, d1', ..., dN') with d_i' = d_i * up / down.

Raises:

Type	Description
ValueError	If up * d_i is not divisible by down for a spatial size d_i.