Neural solver implementation

[pietrocipolla]

Hi!

I want to implement the OT neural solver, but I am a newbie to JAX. In particular, I am interested in the ott.neural.methods.neuraldual.W2NeuralDual function. In the tutorial, it is not clear how to define the data loader class for the neural solvers. In my problem, I create a sample using the following function:

import torch
from torch.distributions.multivariate_normal import MultivariateNormal

def gaussian_model(N, M, device = "cpu"):
    # Create the input tensor
    Sigma = torch.tensor([[1.0, 0.5, 0.5],
                          [0.5, 1.0, 0.5],
                          [0.5, 0.5, 1.0]])
    m_gen = MultivariateNormal(torch.ones(3), Sigma)
    x = m_gen.rsample(sample_shape=torch.Size([N]))
    x = x.to(device)

    # Save for later
    K = x.size(dim=1)

    # Create the output
    A = torch.tensor([[4, -2, 1], [2, 5, -1]], device=device, dtype=torch.float32)
    y = A.mm(x.transpose(1, 0))
    y = y.transpose(1, 0)

    # Order the inputs by column to get the quantiles
    sorted_indices = x.argsort(dim=0, descending=False)

    # Compute the quantiles
    ranks = torch.empty_like(x, dtype=torch.int64, device=device)
    for i in range(K):
        ranks[sorted_indices[:, i], i] = torch.arange(N, device=device)

    # Get the number of elements per partition
    n = torch.arange(start=1, end=N + 2, step=(N / M)).floor()
    n = n.diff().to(dtype=torch.int32)

    # Build identifier for quantiles
    q = torch.repeat_interleave(torch.arange(0, M), n)

    # Define the partitions reordering the 1-M values according to ranks
    partitions = torch.empty((N, K))
    for i in range(K):
        partitions[:, i] = q[ranks[:, i].cpu()]

    # Concatenate the tensors to get the partitions-output tensor
    output = torch.cat((partitions.to(device), y), dim=1)
    output = output[:, [0, 3, 4]]

    return output

Then, I want to compute the distance between the unconditioned distribution of data in columns output[:,1:2] and the distribution obtained by fixing the first variable (output[output[:,0]==m,1:2]). This problem can be easily solved using other OTT solvers such as the Sinkhorn's. However, it is not clear to me how to implement the data loader for the neural solver. I hope that you can help me.