code: def _argmax_helper_all_tree_reductions(
      enc_tensor, dim=None, method="log_reduction"
  ):
      """
      Returns 1 for all elements that have the highest value in the appropriate
      dimension of the tensor. `method` can be one of [`log_reduction`,
      `double_log_reduction`, `accelerated_cascade`].
      `log_reduction`: Uses O(n) comparisons and O(log n) rounds of communication
      `double_log_reduction`: Uses O(n loglog n) comparisons and O(loglog n) rounds
      of communication (Section 2.6.2 in https://folk.idi.ntnu.no/mlh/algkon/jaja.pdf)
      `accelerated_cascade`: Uses O(n) comparisons and O(loglog n) rounds of
      communication. (See Section 2.6.3 of https://folk.idi.ntnu.no/mlh/algkon/jaja.pdf)
      """
      enc_max_vec = _max_helper_all_tree_reductions(
          enc_tensor, dim=dim, method=method
      )
      # reshape back to the original size
      enc_max_vec_orig = enc_max_vec
      if dim is not None:
          enc_max_vec_orig = enc_max_vec.unsqueeze(dim)
      # compute the one-hot vector over the entire tensor
      enc_one_hot_vec = enc_tensor.eq(enc_max_vec_orig)
      return enc_one_hot_vec, enc_max_vec
  
file_path: CrypTen/crypten/common/functions/maximum.py
imports: import math
  import crypten
  import torch
  from crypten.config import cfg
label_desc: Use image node_id 12 for calling a global function with example usage: _argmax_helper_all_tree_reductions(enc_tensor, dim, method) and returns: enc_one_hot_vec,  enc_max_vec