cot.assignment_torch

cot.assignment_torch(W, C, eps, device, seed=1)[source]

This function sloves an additive approximation of the assignment problem between two discrete distributions and returns assignment plan and cost. This function is a torch implementation version of the parallelizable combinatorial algorithm [2] for assignment problem.

Parameters:

  • W (tensor, shape (n, n)) – A n by n cost matrix, each W(i,j) represents the cost between i-th type b and j-th type vertex from the samples of supply and demand distribution.

  • C (tensor) – The scale of cost metric.

  • eps (tensor) – The additive error of optimal transport distance, the value of \(\epsilon\) in paper [2].

  • device (torch.device) – The device where the computation will be executed. (e.g. torch.device(‘cuda:0’) for GPU)

Returns:

  • Mb (tensor, shape (n,)) – A 1 by n array, each i represents the index of type a vertex (j-th column in cost matrix, i.e. W[:,j]) that is assigned with ith type b vertex.

  • yA (tensor, shape (n,)) – A 1 by n array, each i represents the final dual value of ith type a vertex.

  • yB (tensor, shape (n,)) – A 1 by n array, each i represents the final dual value of ith type b vertex.

  • total_cost (float) – Total cost of the final assignment.

References