Sinkhorn implementation not linear w.r.t scaling

[chrzzl]

Very beautiful notebook! However, I believe your implementation of the Sinkhorn distance is faulty:

The formal definition of the Sinkhorn Distance (SD) implies that the SD between two translated point clouds A and A+t is |t| and (2) It is linear wrt to scaling, i.e. SD( sA, sB) = s * SD(A, B) for any two point clouds A and B and scaling param s.
These two conditions are not met in your implementation, however they do hold if I use the SD implementation from the specialized Optimal Transport package “POT” for Python. See below for a test snippet.
Could you please double-check?

import ot

N = 100 # Number nodes
T = 20 # Translation distance
S = 0.5 # Scaling factor

# Generate two 3D point clouds A and B := A + t, where translation vector t has length T.
A = np.random.rand(N, 3)
B = [coord + np.array([T,0,0]) for coord in coords]
# Scale both pointclouds by S
SA = [coord * S for coord in coords]
SB = [coord * S for coord in translated_coords]

def to_array(coords):
    return np.array([coord for coord in coords])

def SinkhornDistanceByYOU(A, B):
    A = to_array(A)
    B = to_array(B)
    A = torch.FloatTensor(A)
    B = torch.FloatTensor(B)
    sinkhorn_distance = SinkhornDistance(eps=1e-7, max_iter=100, reduction='mean')
    cost, pi, C = sinkhorn_distance(A, B)
    return cost

def SinkhornDistanceByPOT(A, B):
    A = to_array(A)
    B = to_array(B)
    # Compute pairwise distance matrix
    M = ot.dist(A, B, metric='euclidean')
    # Uniform weights (assume each point contributes equally)
    w_A = np.ones(len(A)) / len(A)
    w_B = np.ones(len(B)) / len(B)
    # Compute Sinkhorn Distance (Regularized Optimal Transport)
    lambda_reg = 0.05  # Regularization parameter
    S_dist = ot.sinkhorn2(w_A, w_B, M, lambda_reg)
    return S_dist