ott.solvers.nn.neuraldual.NeuralDualSolver.setup

NeuralDualSolver.setup(rng, neural_f, neural_g, input_dim, optimizer_f, optimizer_g)

Setup all components required to train the network.

Return type

None

Parameters

• rng (jax.Array) –

• neural_f (ott.solvers.nn.icnn.ICNN) –

• neural_g (ott.solvers.nn.icnn.ICNN) –

• input_dim (int) –

• optimizer_f (Union[jax.Array, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) –

• optimizer_g (Union[jax.Array, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) –