def log_laplace_loss(x, mu, log_sigma):
  """Computes the log laplace loss as according to ...

  equation (2) in the paper https://arxiv.org/pdf/2102.12092.pdf


  Standard laplace distribution:
  p(x1) = exp(-|x1 - mu| / sigma) / (2 * sigma), -inf < x < +inf.

  Let x1 = logit(x) = log(x / (1 - x))
  p(x) = exp(-|logit(x) - mu| / sigma) / (2 * sigma * x * (1 - x))

  loss = -log(p(x))
       = |logit(x) - mu| / sigma + log(2 * sigma * x * (1 - x))
       = |logit(x) - mu| / sigma + log(sigma) + log(2 * x * (1 - x))

  Args:
    x: tensor of shape [b, ...]. x is expected to be in range [0, 1). Typically,
      x is the groundtruth of a sampled img.
    mu: mean of the log-laplace distribution.
    log_sigma: log of sigma value above. mu and log_sigma are typically
      predictions from a model.

  Returns:
    (loss, neg_logl): a pair of tensors of shape [b, ...]. neg_logl is the
    negative log likelihood of x given p(x|mu, sigma).
  """
  # Cap x to be within [epsilon, 1.0 - epsilon].
  epsilon = 1e-6
  x = jnp.clip(x, epsilon, 1.0 - epsilon)

  logit_x = jnp.log(x / (1.0 - x))
  sigma = jnp.exp(log_sigma)
  loss = jnp.abs(logit_x - mu) / sigma + log_sigma
  # negative log likelihood needs to add in the remaining constant.
  neg_logl = loss + jnp.log(2.0 * x * (1.0 - x))
  return loss, neg_logl
Computes the log laplace loss as according to ...

  equation (2) in the paper https://arxiv.org/pdf/2102.12092.pdf


  Standard laplace distribution:
  p(x1) = exp(-|x1 - mu| / sigma) / (2 * sigma), -inf < x < +inf.

  Let x1 = logit(x) = log(x / (1 - x))
  p(x) = exp(-|logit(x) - mu| / sigma) / (2 * sigma * x * (1 - x))

  loss = -log(p(x))
       = |logit(x) - mu| / sigma + log(2 * sigma * x * (1 - x))
       = |logit(x) - mu| / sigma + log(sigma) + log(2 * x * (1 - x))

  Args:
    x: tensor of shape [b, ...]. x is expected to be in range [0, 1). Typically,
      x is the groundtruth of a sampled img.
    mu: mean of the log-laplace distribution.
    log_sigma: log of sigma value above. mu and log_sigma are typically
      predictions from a model.

  Returns:
    (loss, neg_logl): a pair of tensors of shape [b, ...]. neg_logl is the
    negative log likelihood of x given p(x|mu, sigma).
  
videogvt/train_lib/losses.py