Skip to content

Embeddings¤

Input feature maps for coordinate-based learning.

Note

Key notes:

  • Random Fourier features approximate stationary kernels via \(\phi(x)=[\cos(Bx),\sin(Bx)]\).

phydrax.nn.RandomFourierFeatureEmbeddings ¤

Random Fourier feature embedding.

Samples a (possibly multi-block) Gaussian matrix \(B\) and returns

\[ \phi(x)=\big[\cos(Bx),\ \sin(Bx)\big]. \]
__init__(*, in_size: typing.Union[int, collections.abc.Sequence[int], typing.Literal['scalar']], out_size: int = 32, mu: ArrayLike | collections.abc.Sequence[ArrayLike] = 0.0, sigma: ArrayLike | collections.abc.Sequence[ArrayLike] = 1.0, trainable: bool = False, key: Key[Array, ''] = jr.key(0)) ¤

Arguments:

  • in_size: Input value size. The input is flattened to a vector.
  • out_size: Output feature size (must be even; includes cos and sin parts).
  • mu: Mean for Gaussian \(B\) blocks (scalar or sequence for multiscale).
  • sigma: Standard deviation for Gaussian \(B\) blocks (scalar or sequence for multiscale).
  • trainable: If True, learns \(B\); otherwise stops gradients through \(B\).
  • key: PRNG key.
__call__(x: Array, /, *, key: Key[Array, ''] = jr.key(0)) -> Array ¤