code: def _get_einsum_bias_adjoint_reduction_axes(
      equation: str,
      bias_axes: str,
      einsum_rank: int,
  ) -> list[int]:
      """Computes axes related to the per-example adjoint of the einsum bias op.
  
      To describe the output of this computation, first recall that for each
      example the `EinsumDense` layer performs the following transformation:
      ```
      F(W, bias | X) = Einsum(W, X) + Q(bias)
      ```
      where `W` is a tensor of trainable variables, `bias` is a tensor of rank
      `len(bias_axes)`, `X` is a batch of inputs, and `Q` is a linear broadcast
      operator that roughly corresponds to `Q(bias) ~= tf.broadcast_to(bias, S)` for
      `S := tf.shape(Einsum(W, X))`.
  
      It is straightforward to show that the per-example adjoint of `Q` is given by
      `Q'(Y) := tf.reduce_sum(Y, axes=R)` where `R` contains the broadcasting
      indices. This function returns `R` as an unordered list of `int`s.
  
      Assumptions:
  
        A1. `equation` is one of the following forms:
              C1. `ab,bc->ac`
              C2. `...ab,bc->...ac`
              C3. `ab...,bc->ac...`
  
        A2. The first character in the substring `a` (or `...a` in C2)
            in assumption A1 corresponds to the batch dimension.
  
        A3. The characters in `bias_axes` must be subset of the non-batch dimension
            characters in the substring `ac` (or `...ac` in C2) in
            assumption A1.
  
        A4. `einsum_rank` is the length of the substring `ac` (or `...ac` in C2) in
            assumption A1. This includes the batch dimension.
  
      Examples:
  
        1. equation = 'ab,bc->ac', bias_axes = 'c', einsum_rank = 2 -> []
        2. equation = 'ab,bce->ace', bias_axes = 'ce', einsum_rank = 3, -> []
        3. equation = 'ab,bce->ace', bias_axes = 'c', einsum_rank = 3, -> [2]
        4. equation = 'ab,bce->ace', bias_axes = 'e', einsum_rank = 3, -> [1]
        5. equation = 'ab,bced->aced', bias_axes = 'ced', einsum_rank = 4 -> []
        6. equation = 'ab,bced->aced', bias_axes = 'ce', einsum_rank = 4, -> [3],
        7. equation = 'ab,bced->aced', bias_axes = 'c', einsum_rank = 4, -> [2, 3]
        8. equation = '...ab,bce->...ace', bias_axes = 'c', einsum_rank = 4
           -> [1, 3]
        9. equation = '...ab,bce->...ace', bias_axes = 'c', einsum_rank = 10
           -> [1, 2, 3, 4, 5, 6, 7, 9]
        10. equation = 'ab...,bce->ace...', bias_axes = 'e', einsum_rank = 4
           -> [1, 3]
  
      Args:
        equation: The einsum equation `string`.
        bias_axes: A substring of the output part of `equation` specifying which
          axes a bias `tf.Tensor` is added to.
        einsum_rank: The rank of the tensor that the per-example adjoint operator is
          being applied to.
  
      Returns:
        A list of `int` containing axes in the `input` corresponding to
        `input_rank`. Each `int` is at most `input_rank-1` and excludes zero.
  
      Raises:
        ValueError: If `equation` is not a valid einsum equation in the context of
          the `tf.keras.layers.EinsumDense` layer.
      """
      reduction_axes = []
      bias_char_set = set(bias_axes)
      equation_type, (_, _, ac_str) = _parse_einsum_equation(equation)
      # Do not allow the bias axes to be the batch axis, since we want the adjoint
      # of the bias broadcast op to apply the same operation to all examples in a
      # batch.
      if (
          equation_type != EquationType.LEFT_ELLIPSES
          and ac_str[0] in bias_axes
      ):
          raise ValueError(
              f"Bias axis '{bias_axes}' cannot also be the batch axis."
          )
      # If `equation` of the form `...ab,bc->...ac`, i.e., case C2, we do a
      # right to left traversal; the other cases do a left to right traversal.
      input_indices = range(einsum_rank)
      traversal_zip = (
          itertools.zip_longest(
              reversed(input_indices), reversed(ac_str)
          )
          if equation_type == EquationType.LEFT_ELLIPSES
          else itertools.zip_longest(input_indices, ac_str)
      )
      # Traverse the output part of `equation` and add an index to the output if
      # the corresponding `char` in the `ac` part is NOT in `bias_axes` and the
      # index is not zero (batch dimension). Add all indices except index zero in
      # the `...` part of the output substring (if present).
      for idx, output_char in traversal_zip:
          # Exclude the batch dimension (idx == 0), since we want the per-example
          # adjoint.
          if idx != 0:
              if output_char is not None and bias_char_set:
                  if output_char not in bias_char_set:
                      reduction_axes.append(idx)
                  else:
                      bias_char_set.remove(output_char)
              else:
                  reduction_axes.append(idx)
      return reduction_axes
  
file_path: privacy/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/einsum_utils.py
imports: import enum
  import itertools
  import os
  import re
  from typing import Optional
  import tensorflow
  from tensorflow_privacy.privacy.fast_gradient_clipping import common_manip_utils
label_desc: Use image node_id 6 for calling a global function with example usage: _get_einsum_bias_adjoint_reduction_axes(equation, bias_axes, einsum_rank) and returns: reduction_axes