code: def _reshape_einsum_inputs(
      input_tensor: tf.Tensor,
      equation: str,
  ) -> tf.Tensor:
      """Converts an input tensor of arbitrary rank to a batched matrix tensor.
  
      Args:
        input_tensor: A `tf.Tensor` corresponding to the first input of the einsum
          equation.
        equation: The einsum equation `string`.
  
      Returns:
        A rank-3 `tf.Tensor` representing a batch of rank-2 matrices with the same
        number of rows and columns. The output dimensions, in order, are:
        ```
        (num_batches, num_rows, num_columns)
        ```
        When `input_tensor` is a rank-2 `tf.Tensor`, the number of output rows is 1
        and the number of output columns is the second dimension of the input.
  
      Raises:
        ValueError: If `equation` is not a valid einsum equation in the context of
          the `tf.keras.layers.EinsumDense` layer.
      """
      # Find the components `ab`, `bc`, and `ac` given that `equation` can only be
      # one of the following mutually exclusive forms:
      #
      #   C1. ab,bc->ac,
      #   C2. ...ab,bc->...ac
      #   C3. ab...,bc->ac...
      #
      # NOTE: `a`, `b`, and `c` are (possibly) also substrings.
  
      # Compute the first index of the `b` part of the `ab` component.
      input_shape = input_tensor.shape
      input_len = len(input_shape)
      equation_type, (ab_str, bc_str, ac_str) = _parse_einsum_equation(
          equation
      )
      if equation_type == EquationType.LEFT_ELLIPSES:
          # In case C2, the `a` part of this component can be empty, so we have no
          # choice but to compare the `c` part of `ac` with the `bc` component.
          c_len = 0
          for s1, s2 in itertools.zip_longest(
              reversed(bc_str), reversed(ac_str)
          ):
              if s1 == s2:
                  c_len += 1
              else:
                  break
          b_len = len(bc_str) - c_len
          b_idx = input_len - b_len
      else:
          # For the other cases, we simply compare `ab` with `ac` to get the length
          # of the `a` component, i.e., the first index of `b`.
          b_idx = 0
          for s1, s2 in itertools.zip_longest(ab_str, ac_str):
              if s1 == s2:
                  b_idx += 1
              else:
                  break
      # Prepare `input_tensor` for reshaping and get the pivot index of the prepped
      # tensor. Note that case C3 requires a transpose to ensure that matrix
      # multiplication is performed by the caller.
      if equation_type == EquationType.RIGHT_ELLIPSES:
          ellipses_idx = len(ab_str)
          # Convert `ab...` to `a...b`.
          new_ordering = (
              list(range(0, b_idx))
              + list(range(ellipses_idx, input_len))
              + list(range(b_idx, ellipses_idx))
          )
          input_tensor = tf.transpose(input_tensor, perm=new_ordering)
          ellipses_len = input_len - ellipses_idx
          pivot_idx = b_idx + ellipses_len
      else:
          pivot_idx = b_idx
      # The output tensor is a batched set of matrices, split at the pivot index
      # of the previously prepped tensor.
      input_shape = tf.shape(input_tensor)
      batch_size = input_shape[0]
      num_rows = tf.reduce_prod(input_shape[1:pivot_idx])
      num_columns = tf.reduce_prod(input_shape[pivot_idx:])
      return tf.reshape(
          input_tensor, shape=[batch_size, num_rows, num_columns]
      )
  
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 4 for calling a global function with example usage: _reshape_einsum_inputs(input_tensor, equation) and returns: tf