Skip to content

CNO Blocks

CNO Blocks¤

neojax.nn.cno_blocks.CNOConvBlock ¤

Convolution - batch normalization - filtered activation block of the CNO.

Implements the general n-dimensional convolution operator \(\mathcal{K}_l\) followed by the alias-free activation \(\Sigma_{s, s'}\) of the CNO reference, i.e. the CNO I-, U- and D-blocks. The resampling is part of the activation, see FilteredLReLU.

The filters of the activation are designed as in the reference from the resolution ratio s' / s of the block type: With the input sampling rate \(s\) and the output sampling rate \(s'\) the activation is evaluated at \(s_{\text{act}} = 2 \max(s, s')\). The up-/downsampling filters have cutoffs \(s / c\) and \(s' / c\), with the cutoff denominator \(c\) (cutoff_den), and half-widths \(m s - s / c\) and \(m s' - s' / c\), with the half-width multiplier \(m\) (half_width_mult). This only depends on the ratio of \(s\) and \(s'\), so a block works for any input resolution.

Parameters:

Name Type Description Default
key PRNGKeyArray

PRNG key for parameter initialization.

 required
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
ndim int

Number of spatial dimensions.

 required
block_type Literal['invariant', 'upsample', 'downsample']

The resampling of the block. Must be one of "invariant" (same resolution), "upsample" (twice the resolution) or "downsample" (half the resolution). Default is "invariant".

 'invariant'
conv_kernel int

Odd size of the convolution kernel along each dimension. Default is 3.

 3
batch_norm bool

Whether to apply batch normalization after the convolution, before the activation, as in the reference. If False, there is no normalization. Default is True. Batch normalization makes the block stateful, see below.

 True
batch_axis_name str

Name of the jax.vmap axis the batch statistics are computed over if batch_norm=True. Default is "batch".

 'batch'
batch_norm_momentum float

Weight of the running statistics (not of the batch statistics) in their update if batch_norm=True. Default is 0.9, which corresponds to PyTorch's momentum=0.1.

 0.9
cutoff_den float

Denominator of the cutoff frequency of the filters, in (2, inf). Default is 2.0001, i.e. slightly above critical sampling.

 2.0001
half_width_mult float

Multiplier of the half-width of the transition band of the filters. Default is 0.8.

 0.8
filter_size int

Number of filter taps per phase of the resampling, i.e. the filters have filter_size * factor taps. Default is 6.

 6
pad_mode Literal['wrap', 'reflect', 'constant']

Padding mode of the convolution and the activation. Use "wrap" for a periodic domain, otherwise "reflect" or "constant" (zero padding, as in the reference). Default is "constant".

 'constant'
negative_slope float

Negative slope of the leaky ReLU. Default is 0.2.

 0.2

Raises:

Type Description
ValueError

If conv_kernel is not odd.

Stateful batch normalization

The block is an eqx.nn.StatefulLayer. It is only stateful if batch_norm=True (see is_stateful). In this case, it has to be called as out, state = block(x, state) inside of a jax.vmap with axis_name=batch_axis_name, with the state returned by eqx.nn.make_with_state, see Equinox docs. Otherwise, it is called as out = block(x).

Internal Attributes

These fields store the internal layers state (and weights).

  • convolution (eqx.nn.Conv): The convolution layer.
  • batch_norm (eqx.nn.BatchNorm | None): The batch normalization layer or None.
  • activation (FilteredLReLU): The filtered activation with resampling.

Examples:

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from neojax.nn import CNOConvBlock

key = jr.key(0)
block = CNOConvBlock(
    key=key,
    in_channels=3,
    out_channels=8,
    ndim=2,
    block_type="downsample",
    batch_norm=False,
)
x = jnp.ones((3, 32, 32))
# out.shape = (8, 16, 16)
out = block(x)

# Batch normalization: stateful block
stateful_block, state = eqx.nn.make_with_state(CNOConvBlock)(
    key=key, in_channels=3, out_channels=8, ndim=2, batch_norm=True
)
xs = jnp.ones((4, 3, 32, 32))
out, state = jax.vmap(
    stateful_block, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
)(xs, state)
is_stateful() -> bool ¤

Whether the block is stateful, i.e. uses batch normalization.

Returns:

Type Description
bool

True if the block has to be called as block(x, state) and returns (out, state), False if it is called as block(x).
__call__(x: Inexact[Array, 'in_c ...'], state: equinox.nn._stateful.State | None = None, *, key: PRNGKeyArray | None = None, inference: bool | None = None) -> Inexact[Array, 'out_c ...'] | tuple[Inexact[Array, 'out_c ...'], equinox.nn._stateful.State] ¤

Forward pass of the CNO convolution block.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array shaped (in_c, d1, ..., dN).

 required
state State | None

The state holding the running statistics of the batch normalization. Required if (and only if) is_stateful().

 None
key PRNGKeyArray | None

Unused. Accepted for compatibility with eqx.nn.Sequential.

 None
inference bool | None

Whether to use the running statistics of the batch normalization instead of the batch statistics. If None, the value of the BatchNorm is used (see eqx.nn.inference_mode). Ignored if the block is not stateful.

 None

Returns:

Type Description
Inexact[Array, 'out_c ...'] | tuple[Inexact[Array, 'out_c ...'], State]

Output array shaped (out_c, d1', ..., dN'), with the resolution scaled according to the block_type. If is_stateful(), the output and the updated state.

Raises:

Type Description
ValueError

If a state is required, but not given or vice versa.

neojax.nn.cno_blocks.CNOResidualBlock ¤

Residual block of the CNO.

Computes x + bn(conv(act(bn(conv(x))))) with resolution preserving convolutions, optional batch normalizations bn and the filtered activation (an I-block, see CNOConvBlock).

Parameters:

Name Type Description Default
key PRNGKeyArray

PRNG key for parameter initialization.

 required
channels int

Number of input and output channels.

 required
ndim int

Number of spatial dimensions.

 required
conv_kernel int

Odd size of the convolution kernel along each dimension. Default is 3.

 3
batch_norm bool

Whether to apply batch normalization after each convolution, as in the reference. If False, there is no normalization. Default is True. Batch normalization makes the block stateful, see below.

 True
batch_axis_name str

Name of the jax.vmap axis the batch statistics are computed over if batch_norm=True. Default is "batch".

 'batch'
batch_norm_momentum float

Weight of the running statistics (not of the batch statistics) in their update if batch_norm=True. Default is 0.9, which corresponds to PyTorch's momentum=0.1.

 0.9
cutoff_den float

Denominator of the cutoff frequency of the filters. See CNOConvBlock. Default is 2.0001.

 2.0001
half_width_mult float

Multiplier of the half-width of the transition band of the filters. See CNOConvBlock. Default is 0.8.

 0.8
filter_size int

Number of filter taps per phase of the resampling. See CNOConvBlock. Default is 6.

 6
pad_mode Literal['wrap', 'reflect', 'constant']

Padding mode of the convolutions and the activation. Use "wrap" for a periodic domain, otherwise "reflect" or "constant" (zero padding, as in the reference). Default is "constant".

 'constant'
negative_slope float

Negative slope of the leaky ReLU. Default is 0.2.

 0.2

Raises:

Type Description
ValueError

If conv_kernel is not odd.

Stateful batch normalization

The block is an eqx.nn.StatefulLayer. It is only stateful if batch_norm=True (see is_stateful). See CNOConvBlock for how to call stateful blocks.

Internal Attributes

These fields store the internal layers state (and weights).

  • block (CNOConvBlock): First convolution, batch normalization and activation.
  • convolution (eqx.nn.Conv): The second convolution layer.
  • batch_norm (eqx.nn.BatchNorm | None): The second batch normalization layer or None.

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import CNOResidualBlock

key = jr.key(0)
block = CNOResidualBlock(key=key, channels=8, ndim=2, batch_norm=False)
x = jnp.ones((8, 32, 32))
out = block(x)
is_stateful() -> bool ¤

Whether the block is stateful, i.e. uses batch normalization.

Returns:

Type Description
bool

True if the block has to be called as block(x, state) and returns (out, state), False if it is called as block(x).
__call__(x: Inexact[Array, 'c ...'], state: equinox.nn._stateful.State | None = None, *, key: PRNGKeyArray | None = None, inference: bool | None = None) -> Inexact[Array, 'c ...'] | tuple[Inexact[Array, 'c ...'], equinox.nn._stateful.State] ¤

Forward pass of the residual block.

Parameters:

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

Input array shaped (c, d1, ..., dN).

 required
state State | None

The state holding the running statistics of the batch normalizations. Required if (and only if) is_stateful().

 None
key PRNGKeyArray | None

Unused. Accepted for compatibility with eqx.nn.Sequential.

 None
inference bool | None

Whether to use the running statistics of the batch normalizations instead of the batch statistics. If None, the value of the BatchNorm is used (see eqx.nn.inference_mode). Ignored if the block is not stateful.

 None

Returns:

Type Description
Inexact[Array, 'c ...'] | tuple[Inexact[Array, 'c ...'], State]

Output array with the same shape as the input. If is_stateful(), the output and the updated state.

Raises:

Type Description
ValueError

If a state is required, but not given or vice versa.

neojax.nn.cno_blocks.CNOLiftProjectBlock ¤

Lifting or projection block of the CNO.

Applies an un-normalized, CNOConvBlock that maps the input to latent_channels followed by a final convolution to out_channels. Used both to lift the inputs and to project the latent features to the outputs. As in the reference, there is no batch normalization (stateless).

Parameters:

Name Type Description Default
key PRNGKeyArray

PRNG key for parameter initialization.

 required
in_channels int

Number of input channels.

 required
out_channels int

Number of output channels.

 required
ndim int

Number of spatial dimensions.

 required
latent_channels int

Number of channels between the two layers. Defaults to 64.

 64
conv_kernel int

Odd size of the convolution kernel along each dimension. Default is 3.

 3
cutoff_den float

Denominator of the cutoff frequency of the filters. See CNOConvBlock. Default is 2.0001.

 2.0001
half_width_mult float

Multiplier of the half-width of the transition band of the filters. See CNOConvBlock. Default is 0.8.

 0.8
filter_size int

Number of filter taps per phase of the resampling. See CNOConvBlock. Default is 6.

 6
pad_mode Literal['wrap', 'reflect', 'constant']

Padding mode of the convolutions and the activation. Use "wrap" for a periodic domain, otherwise "reflect" or "constant" (zero padding, as in the reference). Default is "constant".

 'constant'
negative_slope float

Negative slope of the leaky ReLU. Default is 0.2.

 0.2

Raises:

Type Description
ValueError

If conv_kernel is not odd.
Internal Attributes

These fields store the internal layers state (and weights).

  • block (CNOConvBlock): Convolution and filtered activation, without batch normalization.
  • convolution (eqx.nn.Conv): The final convolution layer.

Examples:

import jax.numpy as jnp
import jax.random as jr
from neojax.nn import CNOLiftProjectBlock

key = jr.key(0)
lift = CNOLiftProjectBlock(key=key, in_channels=3, out_channels=8, ndim=2)
x = jnp.ones((3, 32, 32))
# out.shape = (8, 32, 32)
out = lift(x)
__call__(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...'] ¤

Forward pass of the lift/project block.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array shaped (in_c, d1, ..., dN).

 required

Returns:

Type Description
Inexact[Array, 'out_c ...']

Output array shaped (out_c, d1, ..., dN).