code: def _extract_diagonals(formula_a, a):
      """Perform diagonal extraction in einsum formula.
  
      Args:
        formula_a: a string of indices compatible with the given tensor
        a: a tensor compatible with the given indices. For example
          consider the case
            `formula_a == ijik` and `a.shape == [2, 3, 2, 4]`.
          These are compatible because the tensor has rank 4 and
          the string has 4 index letters. Additionally the first and
          third index letters are the same corresponding to the first
          and third dimensions of the tensor shape being equal.
  
      Returns:
        If there are any repeated indices in the index string then
        these correspond to entries in `a` along a generalized
        diagonal. For example the 'ii' entries in a 3x3 matrix `m`
        correspond to the 3 diagonal entries of `m`.
        `_extract_diagonals('ii', m)` returns a pair consisting of
        a 3-vector representing the diagonal of `m` as well as the
        index string `i` indicating that the extraction was
        equivalent to computing `einsum('ii->i', m)`.
        If we have an einsum formula like 'ijik,j,k->jk'
        then we can reduce it to an einsum of the form
        'ijk,j,k->jk', removing the repeated indices,
        by applying `_extract_diagonals` to the corresponding tensor.
        `_extract_diagonals` also returns the string 'ijk'.
      """
  
      extracted_a = ""
      for i in formula_a:
          if i not in extracted_a:
              extracted_a += i
      if formula_a != extracted_a:
          # It is safe to use ordinary `einsum` here as no summations
          # are performed.
          a = tf.einsum("{}->{}".format(formula_a, extracted_a), a)
          formula_a = extracted_a
  
      return formula_a, a
  
file_path: probability/tensorflow_probability/python/experimental/marginalize/logeinsumexp.py
imports: import opt_einsum
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 4 for calling a global function with example usage: _extract_diagonals(formula_a, a) and returns: formula_a,  a