code: 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.
  
      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 = jnp.linalg.svd(jacobian, compute_uv=False)
          log_svals = jnp.log(jnp.maximum(svals, eps))
          sq_residual = jnp.sum(log_svals**2, axis=-1)
      elif loss_type == "svals":
          svals = jnp.linalg.svd(jacobian, compute_uv=False)
          sq_residual = jnp.sum((svals - 1.0) ** 2, axis=-1)
      elif loss_type == "jtj":
          jtj = jacobian @ jacobian.T
          sq_residual = ((jtj - jnp.eye(3)) ** 2).sum() / 4.0
      elif loss_type == "div":
          div = utils.jacobian_to_div(jacobian)
          sq_residual = div**2
      elif loss_type == "det":
          det = jnp.linalg.det(jacobian)
          sq_residual = (det - 1.0) ** 2
      elif loss_type == "log_det":
          det = jnp.linalg.det(jacobian)
          sq_residual = jnp.log(jnp.maximum(det, eps)) ** 2
      elif loss_type == "nr":
          rot = nearest_rotation_svd(jacobian)
          sq_residual = jnp.sum((jacobian - rot) ** 2)
      else:
          raise NotImplementedError(
              f"Unknown elastic loss type {loss_type!r}"
          )
      residual = jnp.sqrt(sq_residual)
      loss = utils.general_loss_with_squared_residual(
          sq_residual, alpha=-2.0, scale=0.03
      )
      return loss, residual
  
imports: import functools
  from typing import Any
  from typing import Callable
  from typing import Dict
  from absl import logging
  from flax import struct
  from flax.training import checkpoints
  import jax
  from jax import lax
  from jax import numpy
  from jax import random
  from jax import vmap
  from nerfies import model_utils
  from nerfies import models
  from nerfies import utils
label_desc: Use image node_id 3 for calling a global function with example usage: compute_elastic_loss(jacobian, eps, loss_type) and returns: loss,  residual