code: def scatter_argmax(
      src: Tensor,
      index: Tensor,
      dim: int = 0,
      dim_size: Optional[int] = None,
  ) -> Tensor:
      if (
          torch_geometric.typing.WITH_TORCH_SCATTER
          and not is_compiling()
      ):
          out = torch_scatter.scatter_max(
              src, index, dim=dim, dim_size=dim_size
          )
          return out[1]
  
      # Only implemented under certain conditions for now :(
      assert src.dim() == 1 and index.dim() == 1
      assert dim == 0 or dim == -1
      assert src.numel() == index.numel()
  
      if dim_size is None:
          dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
  
      if torch_geometric.typing.WITH_PT112:
          res = src.new_empty(dim_size)
          res.scatter_reduce_(
              0, index, src.detach(), reduce="amax", include_self=False
          )
      elif torch_geometric.typing.WITH_PT111:
          res = torch.scatter_reduce(
              src.detach(),
              0,
              index,
              reduce="amax",
              output_size=dim_size,
          )  # type: ignore
      else:
          raise ValueError("'scatter_argmax' requires PyTorch >= 1.11")
  
      out = index.new_full((dim_size,), fill_value=dim_size - 1)
      nonzero = (src == res[index]).nonzero().view(-1)
      out[index[nonzero]] = nonzero
  
      return out
  
file_path: pytorch_geometric/torch_geometric/utils/_scatter.py
imports: from typing import Optional
  import torch
  from torch import Tensor
  import torch_geometric.typing
  from torch_geometric import is_compiling, warnings
  from torch_geometric.typing import torch_scatter
  from torch_geometric.utils.functions import cumsum
label_desc: Use image node_id 2 for calling a global function with example usage: scatter_argmax(src, index, dim, dim_size) and returns: out, out