code: def _reference_l2_project(src_support, src_probs, dst_support):
      """Multi-axis l2_project, implemented using single-axis l2_project.
  
      This is for testing multiaxis_l2_project's consistency with l2_project,
      when used with multi-axis support vs single-axis support.
  
      Args:
        src_support: Zp in l2_project.
        src_probs: P in l2_project.
        dst_support: Zq in l2_project.
  
      Returns:
        src_probs, projected onto dst_support.
      """
      assert src_support.shape == src_probs.shape
  
      # Remove the batch and value axes, and broadcast the rest to a common shape.
      common_shape = np.broadcast(
          src_support[0, ..., 0], dst_support[..., 0]
      ).shape
  
      # If src_* have fewer internal axes than len(common_shape), insert size-1
      # axes.
      while src_support.ndim - 2 < len(common_shape):
          src_support = src_support[:, None, ...]
  
      src_probs = np.reshape(src_probs, src_support.shape)
  
      # Broadcast args' non-batch, non-value axes to common_shape.
      src_support = np.broadcast_to(
          src_support,
          src_support.shape[:1] + common_shape + src_support.shape[-1:],
      )
      src_probs = np.broadcast_to(src_probs, src_support.shape)
      dst_support = np.broadcast_to(
          dst_support, common_shape + dst_support.shape[-1:]
      )
  
      output_shape = (src_support.shape[0],) + dst_support.shape
  
      # Collapse all but the first (batch) and last (atom) axes.
      src_support = src_support.reshape(
          [src_support.shape[0], -1, src_support.shape[-1]]
      )
      src_probs = src_probs.reshape(
          [src_probs.shape[0], -1, src_probs.shape[-1]]
      )
  
      # Collapse all but the last (atom) axes.
      dst_support = dst_support.reshape([-1, dst_support.shape[-1]])
  
      dst_probs = np.zeros(
          src_support.shape[:1] + dst_support.shape,
          dtype=src_probs.dtype,
      )
  
      # iterate over all supports
      for i in range(src_support.shape[1]):
          s_support = tf.convert_to_tensor(src_support[:, i, :])
          s_probs = tf.convert_to_tensor(src_probs[:, i, :])
          d_support = tf.convert_to_tensor(dst_support[i, :])
          d_probs = distributional.l2_project(
              s_support, s_probs, d_support
          )
          dst_probs[:, i, :] = d_probs.numpy()
  
      return dst_probs.reshape(output_shape)
  
file_path: acme/acme/tf/losses/distributional_test.py
imports: from acme.tf.losses import distributional
  import numpy
  from numpy import testing
  import tensorflow
  from absl.testing import absltest
  from absl.testing import parameterized
label_desc: Use image node_id 1 for calling a global function with example usage: _reference_l2_project(src_support, src_probs, dst_support) and returns: dst_probs