kd.nn.FourierEmbedding

class kauldron.modules.FourierEmbedding(

num_fourier_bases: int,

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>,

name: str | None = None,

)

Bases: flax.linen.module.Module

Apply Fourier position embedding to a grid of coordinates.

num_fourier_bases

The number of Fourier bases to use. The embedding dimensionality is 2 x len(axis) x num_fourier_bases, but the result will be projected to match the given shape.

Type:

int

Returns:

Fourier position embeddings broadcastable to given shape.

num_fourier_bases: int

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: Scope | None = None