code: def pairwise_distance_wrapper(
      sources, targets, weights=1.0, distance_config=None
  ):
      """A wrapper to compute the pairwise distance between `sources` and `targets`.
  
      `distances = weights * distance_config.distance_type(sources, targets)`
  
      This wrapper calculates the weighted distance between `(sources, targets)`
      pairs, and provides an option to return the distance as the sum over the
      difference along the given axis, when vector based distance is needed.
  
      For the usage of `weights` and `reduction`, please refer to `tf.losses`. For
      the usage of `sum_over_axis`, see the following examples:
  
      Given target tensors with shape `[batch_size, features]`, the reduction set to
      `tf.compat.v1.losses.Reduction.MEAN`, and `sum_over_axis` set to the last
      dimension, the weighted average distance of sample pairs will be returned.
      For example: With a distance_config('L2', sum_over_axis=-1), the distance
      between [[1, 1], [2, 2], [0, 2], [5, 5]] and [[1, 1], [0, 2], [4, 4], [1, 4]]
      will be {(0+0) + (4+0) + (16+4) + (16+1)}/4 = 10.25
  
      If `sum_over_axis` is `None`, the weighted average distance of feature pairs
      (instead of sample pairs) will be returned. For example: With a
      distance_config('L2'), the distance between
      [[1, 1], [2, 2], [0, 2], [5, 5]] and [[1, 1], [0, 2], [4, 4], [1, 4]] will be
      {(0+0) + (4+0) + (16+4) + (16+1)}/8 = 5.125
  
      If `transform_fn` is not `None`, the transform function is applied to both
      `sources` and `targets` before computing the distance. For example:
      `distance_config('KL_DIVERGENCE', sum_over_axis=-1, transform_fn='SOFTMAX')`
      treats `sources` and `targets` as logits, and computes the KL-divergence
      between the two probability distributions.
  
      Args:
        sources: `Tensor` of type `float32` or `float64`.
        targets: `Tensor` of the same type and shape as `sources`.
        weights: (optional) `Tensor` whose rank is either 0, or the same as that of
          `targets`, and must be broadcastable to `targets` (i.e., all dimensions
          must be either `1`, or the same as the corresponding distance dimension).
        distance_config: An instance of `nsl.configs.DistanceConfig` that contains
          the following configuration (or hyperparameters) for computing distances:
          (a) `distance_type`: Type of distance function to apply.
          (b) `reduction`: Type of distance reduction. See `tf.losses.Reduction`.
          (c) `sum_over_axis`: (optional) The distance is the sum over the
            difference along the specified axis. Note that if `sum_over_axis` is not
            `None` and the rank of `weights` is non-zero, then the size of `weights`
            along `sum_over_axis` must be 1.
          (d) `transform_fn`: (optional) If set, both `sources` and `targets` will
            be transformed before calculating the distance. If set to 'SOFTMAX', it
            will be performed on the axis specified by 'sum_over_axis', or -1 if the
            axis is not specified. If `None`, the default distance config will be
            used.
  
      Returns:
        Weighted distance scalar `Tensor`. If `reduction` is
          `tf.compat.v1.losses.Reduction.MEAN`, this has the same shape as
          `targets`.
      Raises:
        ValueError: If the shape of targets doesn't match that of sources, or if the
          shape of weights is invalid.
        TypeError: If the distance function gets an unexpected keyword argument.
      """
      if distance_config is None:
          distance_config = configs.DistanceConfig()  # Default configs.
  
      tf.compat.v1.losses.Reduction.validate(distance_config.reduction)
  
      if distance_config.transform_fn is not configs.TransformType.NONE:
          sources = _apply_transform(
              sources,
              distance_config.transform_fn,
              distance_config.sum_over_axis,
          )
          targets = _apply_transform(
              targets,
              distance_config.transform_fn,
              distance_config.sum_over_axis,
          )
  
      sum_over_axis = distance_config.sum_over_axis
      # Validates the `sum_over_axis`
      _assert_valid_axis(sources.get_shape().ndims, sum_over_axis)
      distance_fn = _select_distance_fn(distance_config.distance_type)
      if distance_config.distance_type == configs.DistanceType.COSINE:
          # Cosine distance function assumes input tensors have been unit-normalized
          sources = tf.nn.l2_normalize(sources, axis=sum_over_axis)
          targets = tf.nn.l2_normalize(targets, axis=sum_over_axis)
      if _is_axis_required_in_distance_fn(
          distance_config.distance_type
      ):
          distances = distance_fn(
              labels=sources,
              predictions=targets,
              weights=weights,
              axis=sum_over_axis,
              reduction=distance_config.reduction,
              loss_collection=None,
          )
      else:
          distances = distance_fn(
              labels=sources,
              predictions=targets,
              weights=weights,
              reduction=distance_config.reduction,
              loss_collection=None,
          )
          if sum_over_axis is not None and _is_reduced_by_average(
              distance_config.reduction
          ):
              # The distance is divided by the size of targets tensor, so we need to
              # rescale the distance by multiplying the size of axis. Note, the distance
              # function with `axis` as a required argument (e.g., consine distance)
              # does not need to be rescaled.
              weights = tf.convert_to_tensor(value=weights)
              weights_shape = weights.get_shape().as_list()
              if weights_shape and weights_shape[sum_over_axis] != 1:
                  raise ValueError(
                      "Shape of weights along the axis %d must be 1."
                      % sum_over_axis
                  )
              distances *= sources.shape.dims[sum_over_axis].value
      return distances
  
file_path: neural-structured-learning/neural_structured_learning/lib/distances.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  import tensorflow
label_desc: Use image node_id 10 for calling a global function with example usage: pairwise_distance_wrapper(sources, targets, weights, distance_config) and returns: distances