ott.initializers.nn.initializers.MetaMLP#

class ott.initializers.nn.initializers.MetaMLP(potential_size, num_hidden_units=512, num_hidden_layers=3, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Potential for MetaInitializer.

This provides an MLP \(\hat f_\theta(a, b)\) that maps from the probabilities of the measures to the optimal dual potentials \(f\).

Parameters:

  • potential_size (int) – The dimensionality of \(f\).

  • num_hidden_units (int) – The number of hidden units in each layer.

  • num_hidden_layers (int) – The number of hidden layers.

  • parent (Union[Type[Module], Type[Scope], Type[_Sentinel], None]) –

  • name (Optional[str]) –

Methods

apply(variables, *args[, rngs, method, ...])

Applies a module method to variables and returns output and modified variables.

clone(*[, parent])

Creates a clone of this Module, with optionally updated arguments.

get_variable(col, name[, default])

Retrieves the value of a Variable.

has_rng(name)

Returns true if a PRNGSequence with name name exists.

has_variable(col, name)

Checks if a variable of given collection and name exists in this Module.

init(rngs, *args[, method, mutable, ...])

Initializes a module method with variables and returns modified variables.

init_with_output(rngs, *args[, method, ...])

Initializes a module method with variables and returns output and modified variables.

is_initializing()

Returns True if running under self.init(...) or nn.init(...)().

is_mutable_collection(col)

Returns true if the collection col is mutable.

lazy_init(rngs, *args[, method, mutable])

Initializes a module without computing on an actual input.

make_rng(name)

Returns a new RNG key from a given RNG sequence for this Module.

param(name, init_fn, *init_args[, unbox])

Declares and returns a parameter in this Module.

perturb(name, value[, collection])

Add an zero-value variable ('perturbation') to the intermediate value.

put_variable(col, name, value)

Updates the value of the given variable if it is mutable, or an error otherwise.

setup()

Initializes a Module lazily (similar to a lazy __init__).

sow(col, name, value[, reduce_fn, init_fn])

Stores a value in a collection.

tabulate(rngs, *args[, depth, ...])

Creates a summary of the Module represented as a table.

unbind()

Returns an unbound copy of a Module and its variables.

variable(col, name[, init_fn, unbox])

Declares and returns a variable in this Module.

Attributes

name

num_hidden_layers

num_hidden_units

parent

scope

variables

Returns the variables in this module.

potential_size