In contrast to most methods presented in ott.solvers, which output vectors or matrices, the goal of the ott.neural module is to parameterize optimal transport maps and couplings as neural networks. These neural networks can generalize to new samples, in the sense that they can be conveniently evaluated outside training samples. This module implements layers, models and solvers to estimate such neural networks.


models.ICNN(dim_data, dim_hidden[, ranks, ...])

Input convex neural network (ICNN).

models.MLP(dim_hidden[, is_potential, ...])

A generic, not-convex MLP.

models.MetaInitializer(geom, meta_model[, ...])

Meta OT Initializer with a fixed geometry [Amos et al., 2022].


losses.monge_gap(map_fn, reference_points[, ...])

Monge gap regularizer [Uscidda and Cuturi, 2023].

losses.monge_gap_from_samples(source, target)

Monge gap, instantiated in terms of samples before / after applying map.


layers.PositiveDense(dim_hidden[, ...])

A linear transformation using a matrix with all entries non-negative.

layers.PosDefPotentials(num_potentials[, ...])

\(\frac{1}{2} x^T (A_i A_i^T + \text{Diag}(d_i)) x + b_i^T x^2 + c_i\) potentials.