neuralop.layers.normalization_layers.AdaIN

class neuralop.layers.normalization_layers.AdaIN(embed_dim, in_channels, mlp=None, eps=1e-05)[source]

Adaptive Instance Normalization (AdaIN) layer for style transfer in neural operators.

AdaIN performs instance normalization followed by adaptive scaling and shifting based on an external embedding vector. This allows for style transfer by modulating the output characteristics based on a conditioning signal.

The layer first normalizes the input using instance normalization, then applies learned scaling (weight) and shifting (bias) parameters derived from an embedding vector through a multi-layer perceptron.

Parameters:

embed_dimint

Dimension of the embedding vector used for style conditioning

in_channelsint

Number of input channels to normalize

mlpnn.Module, optional

Multi-layer perceptron that maps embedding to (weight, bias) parameters. Should output 2*in_channels values. If None, a default MLP is created with architecture: Linear(embed_dim, 512) -> GELU -> Linear(512, 2*in_channels)

epsfloat, optional

Small value added to the denominator for numerical stability in normalization. Default is 1e-5.

Methods

forward(x)	Apply adaptive instance normalization to the input tensor.
set_embedding(x)	Set the embedding vector for style conditioning.

set_embedding(x)[source]

Set the embedding vector for style conditioning.

forward(x)[source]

Apply adaptive instance normalization to the input tensor.