import torch


def dist(batch_reprs, eps = 1e-16, squared=False):
	"""
	Efficient function to compute the distance matrix for a matrix A.

	Args:
		batch_reprs:  vector representations
		eps: float, minimal distance/clampling value to ensure no zero values.
	Returns:
		distance_matrix, clamped to ensure no zero values are passed.
	"""
	prod = torch.mm(batch_reprs, batch_reprs.t())
	norm = prod.diag().unsqueeze(1).expand_as(prod)
	res = (norm + norm.t() - 2 * prod).clamp(min = 0)

	if squared:
		return res.clamp(min=eps)
	else:
		return res.clamp(min = eps).sqrt()