ott.solvers.nn#
This module implements the input-convex neural network [Amos et al., 2017] which can be used to solve OT problems between point-clouds. Other simpler alternatives, such as MLP are implemented, essentially borrowed from the flax library.
Neural Dual#
|
Solver for the Wasserstein-2 Kantorovich dual between Euclidean spaces. |
Models#
|
Base class for the neural solver models. |
|
Input convex neural network (ICNN) architecture with initialization. |
|
A non-convex MLP. |
Conjugate Solvers#
|
Holds the results of numerically conjugating a function. |
Abstract conjugate solver class. |
|
Solve for the conjugate using |