RandomFourierKernelND#

class RandomFourierKernelND(
out_dim,
data_dim,
mlp_hidden_dim,
num_layers,
embedding_dim,
omega_0,
L_cache,
use_bias,
nonlinear_cfg,
init_method=None,
)#

Bases: Module

Learned convolutional kernel parametrised via Random Fourier Features and an MLP.

Mathematical form#

The kernel at grid coordinate x is:

k(x) = Linear_out( MLP( phi(x) ) )

where:

  • phi(x) = [cos(W x + b), sin(W x + b)] is the RFF positional embedding (see RandomFourierPositionalEmbeddingND).

  • MLP is a stack of num_layers - 1 fully-connected layers each followed by the nonlinear_cfg activation.

  • Linear_out is a final linear layer that maps to out_dim channels.

The output is a kernel tensor of shape [1, *spatial_dims, out_dim] suitable for passing directly to the FFT convolution primitives in nvsubquadratic.ops (after rearranging to channels-first layout via the consuming CKConvND module).

Hyperparameters controlling bandwidth / smoothness#

  • omega_0: Controls the frequency content of the RFF features. Higher values concentrate the kernel’s spectral energy at higher frequencies, producing a narrower, higher-bandwidth filter. Typical range: 1.0–100.0.

  • mlp_hidden_dim: Width of the hidden MLP layers. Larger values allow more expressive kernel shapes.

  • num_layers: Depth of the MLP. Must be >= 2 (one hidden layer + output layer minimum).

Initialisation#

  • Hidden MLP layers are initialised with the user-supplied init_method if provided; otherwise use PyTorch defaults.

  • The output layer applies Wang initialisation: weights are scaled by sqrt(1 / kernel_volume) where kernel_volume = prod(L_cache_per_axis) (collapses to L_cache**data_dim for an isotropic grid). This normalises the kernel’s initial energy to be independent of grid size.

out_dim#

Number of output channels (kernel depth).

Type:

int

data_dim#

Number of spatial / temporal input dimensions.

Type:

int

mlp_hidden_dim#

Hidden width of the MLP.

Type:

int

num_layers#

Total number of MLP layers (>= 2).

Type:

int

embedding_dim#

RFF embedding dimensionality (must be even).

Type:

int

omega_0#

Bandwidth scaling factor for the positional embedding.

Type:

float

L_cache_per_axis#

Per-axis cache extents (canonical form).

Type:

tuple[int, …]

L_cache#

Original L_cache argument (diagnostics).

Type:

int | Sequence[int]

positional_embedding#

RFF encoder.

Type:

RandomFourierPositionalEmbeddingND

kernel_network#

Hidden MLP layers.

Type:

torch.nn.Sequential

out_linear#

Final projection to out_dim channels.

Type:

torch.nn.Linear

__init__(
out_dim,
data_dim,
mlp_hidden_dim,
num_layers,
embedding_dim,
omega_0,
L_cache,
use_bias,
nonlinear_cfg,
init_method=None,
)#

Initialize the RandomFourierKernelND class.

Parameters:

  • out_dim (int) – Number of output channels for the generated kernel.

  • data_dim (int) – Number of spatial/temporal input dimensions (size of coordinate vector).

  • mlp_hidden_dim (int) – Hidden width of the network.

  • num_layers (int) – Total number of layers including the first and hidden layers (>= 2).

  • embedding_dim (int) – Dimensionality of the positional embeddings.

  • omega_0 (float) – Frequency scaling factor for the positional embeddings.

  • L_cache (int | Sequence[int]) – Per-axis cache extents. Either a scalar int (broadcast to every axis, isotropic grid of size L_cache**data_dim) or a sequence of length data_dim (anisotropic grid of size prod(L_cache)). The Wang init below uses the product of per-axis extents so it stays correct in both cases.

  • use_bias (bool) – Whether to use bias in the network and embedding layers.

  • nonlinear_cfg (LazyConfig) – Configuration for the nonlinear activation function.

  • init_method (Callable[[int], Callable[[Tensor], Tensor]] | None) – Optional initialization method for the kernel network.

forward(seq_lens, conditioning=None)#

Compute the RFF kernel for a given grid of spatial dimensions.

Parameters:

  • seq_lens (tuple[int, ...]) – Per-axis output sequence lengths. Length must equal self.data_dim.

  • conditioning (Tensor | None) – Unused. Accepted for API compatibility with FiLM-enabled kernels (e.g. SIRENKernelND with film_cfg).

Returns:

  • torch.Tensor: Kernel values of shape [1, *spatial_dims, out_dim].

  • torch.Tensor: Coordinate grid of shape [1, *spatial_dims, data_dim].

Return type:

tuple

Parameters: