import torch


def compute_tv_norm(values, losstype = "l2"):
  """Returns TV norm for input values.

  Source: regnerf/internal/math.py

  Args:
    values: [batch, H, W, *]. 3 or more dimensional tensor.
    losstype: l2 or l1

  Returns:
    loss: [batch, H-1, W-1, *]
  """
  v00 = values[:, :-1, :-1]
  v01 = values[:, :-1, 1:]
  v10 = values[:, 1:, :-1]

  if losstype == "l2":
    loss = ((v00 - v01)**2) + ((v00 - v10)**2)
  elif losstype == "l1":
    loss = torch.abs(v00 - v01) + torch.abs(v00 - v10)
  else:
    raise ValueError(f"Unsupported TV losstype {losstype}.")

  return loss