network

Neural network architectures for score-based diffusion models.

This module provides MLP-based score networks with optional Fourier feature embeddings for periodic domains.

MLP

Bases: Module

Simple feed-forward multilayer perceptron producing scalar outputs.

Parameters:

• features (Sequence[int]) –

  Number of units in each hidden layer. A final Dense layer of size 1 is appended automatically.

Returns:

• output ( ndarray ) –

  Array of shape (..., 1), the scalar output for each example.

ScoreMLP

Bases: Module

Time-dependent score network using a simple MLP.

This network concatenates the state x and time embedding t, passes them through an MLP, and sums the final outputs to a scalar score.

Parameters:

• features (Sequence[int]) –

  Hidden layer sizes for the underlying MLP.

Returns:

• score ( ndarray ) –

  Scalar score per example, shape (...,).

ScorePeriodicMLP

Bases: Module

Periodic score network with sine/cosine Fourier features.

Extends ScoreMLP by first embedding x via Fourier features before concatenating with time.

Parameters:

• features (Sequence[int]) –

  Hidden layer sizes for the underlying MLP.

• fourier_features_stop (int) –

  Maximum frequency (inclusive) for Fourier embeddings. Default is 1.

Returns:

• score ( ndarray ) –

  Scalar score per example, shape (...,).

ScoreSymmetricPeriodicMLP

Bases: Module

Symmetric periodic score with cosine-only Fourier features.

Similar to ScorePeriodicMLP but uses only cosine terms for even symmetry.

Parameters:

• features (Sequence[int]) –

  Hidden layer sizes for the underlying MLP.

• fourier_features_stop (int) –

  Maximum frequency (inclusive) for Fourier embeddings. Default is 1.

Returns:

• score ( ndarray ) –

  Scalar score per example, shape (...,).

FourierFeatures

Bases: Module

Generate sine/cosine Fourier features for periodic inputs.

Embeds each dimension of x into multiple frequencies via sine and/or cosine.

Parameters:

• start (int) –

  Starting frequency (inclusive). Default is 1.

• stop (int) –

  Stopping frequency (inclusive). Default is 8.

• step (int) –

  Frequency step size. Default is 1.

• odd (bool) –

  If True, include sine components. Default is True.

• even (bool) –

  If True, include cosine components. Default is True.

Raises:

• ValueError –

  If both odd and even are False.

Returns:

• features ( ndarray ) –

  Flattened array of shape (..., n_dims * n_freqs * n_components), where n_components is 1 or 2 depending on odd/even.