code: def one_hot_add(inputs, shift):
      """Performs (inputs + shift) % vocab_size in the one-hot space.
      Args:
          inputs: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
          Tensor.
          shift: Tensor of shape `[..., vocab_size]`. Typically a soft/hard one-hot
          Tensor specifying how much to shift the corresponding one-hot vector in
          inputs. Soft values perform a "weighted shift": for example,
          shift=[0.2, 0.3, 0.5] performs a linear combination of 0.2 * shifting by
          zero; 0.3 * shifting by one; and 0.5 * shifting by two.
      Returns:
          Tensor of same shape and dtype as inputs.
      """
      inputs = torch.stack((inputs, torch.zeros_like(inputs)), dim=-1)
      shift = torch.stack((shift, torch.zeros_like(shift)), dim=-1)
      if "torch.fft" not in sys.modules:
          with warnings.catch_warnings(record=True) as w:
              inputs_fft = torch.fft(
                  inputs, 1
              )  # ignore last and first dimension to do batched fft
              shift_fft = torch.fft(shift, 1)
      else:
          inputs_fft = torch.view_as_real(
              torch.fft.fft(torch.view_as_complex(inputs))
          )
          shift_fft = torch.view_as_real(
              torch.fft.fft(torch.view_as_complex(shift))
          )
      result_fft_real = (
          inputs_fft[..., 0] * shift_fft[..., 0]
          - inputs_fft[..., 1] * shift_fft[..., 1]
      )
      result_fft_imag = (
          inputs_fft[..., 0] * shift_fft[..., 1]
          + inputs_fft[..., 1] * shift_fft[..., 0]
      )
      result_fft = torch.stack(
          (result_fft_real, result_fft_imag), dim=-1
      )
      if "torch.fft" not in sys.modules:
          with warnings.catch_warnings(record=True) as w:
              return torch.ifft(result_fft, 1)[
                  ..., 0
              ]  # return only the real part
      else:
          return torch.view_as_real(
              torch.fft.ifft(torch.view_as_complex(result_fft))
          )[..., 0]
  
file_path: DIG/dig/ggraph/method/GraphDF/model/df_utils.py
imports: import sys
  import warnings
  import torch
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: one_hot_add(inputs, shift) and returns: torch, torch