import torch


def sinkhorn(log_alpha, n_iters = 20):
	# torch version
	"""Performs incomplete Sinkhorn normalization to log_alpha.
	By a theorem by Sinkhorn and Knopp [1], a sufficiently well-behaved  matrix
	with positive entries can be turned into a doubly-stochastic matrix
	(i.e. its rows and columns add up to one) via the successive row and column
	normalization.
	-To ensure positivity, the effective input to sinkhorn has to be
	exp(log_alpha) (element wise).
	-However, for stability, sinkhorn works in the log-space. It is only at
	return time that entries are exponentiated.
	[1] Sinkhorn, Richard and Knopp, Paul.
	Concerning nonnegative matrices and doubly stochastic
	matrices. Pacific Journal of Mathematics, 1967
	Args:
	log_alpha: a 2D tensor of shape [N, N]
	n_iters: number of sinkhorn iterations (in practice, as little as 20
		iterations are needed to achieve decent convergence for N~100)
	Returns:
	A 3D tensor of close-to-doubly-stochastic matrices (2D tensors are
		converted to 3D tensors with batch_size equals to 1)
	"""
	n = log_alpha.size()[1]
	log_alpha = log_alpha.view(-1, n, n)

	for i in range(n_iters):
		log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=2, keepdim=True)).view(-1, n, 1)
		log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=1, keepdim=True)).view(-1, 1, n)
	return torch.exp(log_alpha)