ott.problems.nn.dataset.GaussianMixture#

class ott.problems.nn.dataset.GaussianMixture(name, batch_size, init_rng, scale=5.0, variance=0.5)[source]#

A mixture of Gaussians.

Parameters:

  • name (Literal['simple', 'circle', 'square_five', 'square_four']) –

    the name specifying the centers of the mixture components:

    • simple - data clustered in one center,

    • circle - two-dimensional Gaussians arranged on a circle,

    • square_five - two-dimensional Gaussians on a square with one Gaussian in the center, and

    • square_four - two-dimensional Gaussians in the corners of a rectangle

  • batch_size (int) – batch size of the samples

  • init_rng (PRNGKeyArray) – initial PRNG key

  • scale (float) – scale of the individual Gaussian samples

  • variance (float) – the variance of the individual Gaussian samples

Methods

Attributes

scale

variance

name

batch_size

init_rng