code: def dice_loss(
      inputs: jnp.ndarray,
      targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      all_pairs: bool = False,
      eps: float = 1.0,
      interpolation: str = "nearest",
  ) -> jnp.ndarray:
      """Computes the Dice loss given panoptic segmentation logits and targets.
  
      This loss is based on the Dice coefficient (F-1 score). For details, see
      https://arxiv.org/abs/2005.12872 and https://arxiv.org/pdf/1606.04797.pdf.
  
      Args:
        inputs: Predicted mask logits with shape [batch, num_objects, H, W].
        targets: Target masks with shape [batch, num_objects, H, W].
        weights: Array of shape [batch, ...].
        all_pairs: Whether to compute the loss for all object pairs or not.
        eps: Epsilon for numerical stability.
        interpolation: Method to use for upsampling inputs to target size.
  
      Returns:
        If all_pairs == True, returns a [bs, n, m] pairwise matrix, of dice loss.
        If all_pairs == False, returns a [bs, n] matrix of dice loss.
      """
      _, n, h, w = inputs.shape
      b, m, _, _ = targets.shape
  
      # Downsample targets to match prediction:
      # TODO(mjlm): Check if it would be better to upsample predictions.
      # For now, we downsample targets to save memory.
      targets = jax.image.resize(
          targets,
          shape=[b, m, h, w],
          method=interpolation,
          antialias=True,
      )
  
      # TODO(mjlm): Also try softmax instead of sigmoid:
      # As in MaX-DeepLab:
      inputs = jax.nn.sigmoid(inputs)
  
      inputs = jnp.reshape(inputs, [b, n, h * w])
      targets = jnp.reshape(targets, [b, m, h * w])
      if all_pairs:
          numerator = 2 * jnp.einsum("bnp,bkp->bnk", inputs, targets)
          denominator = jnp.sum(
              inputs[:, :, None, :], axis=-1
          ) + jnp.sum(targets[:, None, :, :], axis=-1)
      else:
          assert n == m
          numerator = 2 * jnp.einsum("bnp,bnp->bn", inputs, targets)
          denominator = jnp.sum(inputs + targets, axis=-1)
      loss = 1.0 - (numerator + eps) / (denominator + eps)
  
      if weights is not None:
          loss = apply_weights(loss, weights)
  
      return loss
  
imports: import functools
  from typing import Optional, Any, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 26 for calling a global function with example usage: dice_loss(inputs, targets, weights, all_pairs, eps, interpolation) and returns: loss