Yes I think that’s their particular application in the paper - but it could be more general than that?
The layer wasserstein-loss.jl coded in Mocha, seems to be fairly flexible.
You might have seen this Python optimal transport library,

  
      github.com
  
  
    rflamary/POT/blob/master/ot/bregman.py 180
# -*- coding: utf-8 -*-
"""
Bregman projections for regularized OT
"""

# Author: Remi Flamary <remi.flamary@unice.fr>
#         Nicolas Courty <ncourty@irisa.fr>
#
# License: MIT License

import numpy as np


def sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    u"""
    Solve the entropic regularization optimal transport problem and return the OT matrix

    The function solves the following optimization problem:

    .. math::


  This file has been truncated. show original

  
  
    
    
  
  


I’m guessing that’s what they’ve implemented as their layer? Seems to be a solid piece of theory?