import torch


def metric_fixed_point(cost_matrix, gamma=0.99, eps=1e-7):
  """DP for calculating PSM (approximately).

  Args:
    cost_matrix: DIST matrix where entries at index (i, j) is DIST(x_i, y_j)
    gamma: Metric discount factor.
    eps: Threshold for stopping the fixed point iteration.
  """
  d = torch.zeros_like(cost_matrix)
  def operator(d_cur):
    d_new = 1 * cost_matrix
    discounted_d_cur = gamma * d_cur
    d_new[:-1, :-1] += discounted_d_cur[1:, 1:]
    d_new[:-1, -1] += discounted_d_cur[1:, -1]
    d_new[-1, :-1] += discounted_d_cur[-1, 1:]
    return d_new

  while True:
    d_new = operator(d)
    if torch.sum(torch.abs(d - d_new)) < eps:
      break
    else:
      d = d_new[:]
  return d