transport.matrix.ops.unbalanced_sinkhorn

def unbalanced_sinkhorn(cost: torch.Tensor, eps: float, n_iter: int, rho: float, a: torch.Tensor, b: torch.Tensor, mask: torch.Tensor | None = None, scaling: float | None = None) -> torch.Tensor