code: def _argmax_helper_pairwise(enc_tensor, dim=None):
      """Returns 1 for all elements that have the highest value in the appropriate
      dimension of the tensor. Uses O(n^2) comparisons and a constant number of
      rounds of communication
      """
      dim = -1 if dim is None else dim
      row_length = (
          enc_tensor.size(dim) if enc_tensor.size(dim) > 1 else 2
      )
  
      # Copy each row (length - 1) times to compare to each other row
      a = enc_tensor.expand(row_length - 1, *enc_tensor.size())
  
      # Generate cyclic permutations for each row
      b = mpc.stack(
          [
              enc_tensor.roll(i + 1, dims=dim)
              for i in range(row_length - 1)
          ]
      )
  
      # Use either prod or sum & comparison depending on size
      if row_length - 1 < torch.iinfo(torch.long).bits * 2:
          pairwise_comparisons = a.ge(b)
          result = pairwise_comparisons.prod(0)
      else:
          # Sum of columns with all 1s will have value equal to (length - 1).
          # Using ge() since it is slightly faster than eq()
          pairwise_comparisons = a.ge(b)
          result = pairwise_comparisons.sum(0).ge(row_length - 1)
      return result, None
  
file_path: FATE/python/fate/arch/protocol/mpc/functions/maximum.py
imports: import math
  import torch
  from fate.arch.protocol import mpc
  from fate.arch.protocol.mpc.config import cfg
label_desc: Use image node_id 5 for calling a global function with example usage: _argmax_helper_pairwise(enc_tensor, dim) and returns: result,  None