import torch


def compute_elastic_loss(jacobian, eps=1e-6, loss_type="log_svals"):
  """Compute the elastic regularization loss.

  The loss is given by sum(log(S)^2). This penalizes the singular values
  when they deviate from the identity since log(1) = 0.0,
  where D is the diagonal matrix containing the singular values.
  Based on https://github.com/google/nerfies/blob/main/nerfies/training.py#L71.

  Args:
    jacobian: the Jacobian of the point transformation.
    eps: a small value to prevent taking the log of zero.
    loss_type: which elastic loss type to use.

  Returns:
    The elastic regularization loss.
  """
  if loss_type == "log_svals":
    svals = torch.linalg.svdvals(jacobian)
    log_svals = torch.log(torch.clamp(svals, min=eps, max=None))
    sq_residual = torch.sum(log_svals**2, dim=-1)
  else:
    raise NotImplementedError(f"Unknown elastic loss type {loss_type!r}")
  residual = torch.sqrt(sq_residual)
  loss = sq_residual.mean()
  return loss, residual.mean()