content: import torch
  import torch.nn as nn
  
  
  class FrequencyPositionalEmbedding(nn.Module):
      """The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts
      each feature dimension of `x[..., i]` into:
          [
              sin(x[..., i]),
              sin(f_1*x[..., i]),
              sin(f_2*x[..., i]),
              ...
              sin(f_N * x[..., i]),
              cos(x[..., i]),
              cos(f_1*x[..., i]),
              cos(f_2*x[..., i]),
              ...
              cos(f_N * x[..., i]),
              x[..., i]     # only present if include_input is True.
          ], here f_i is the frequency.
  
      Denote the space is [0 / num_freqs, 1 / num_freqs, 2 / num_freqs, 3 / num_freqs, ..., (num_freqs - 1) / num_freqs].
      If logspace is True, then the frequency f_i is [2^(0 / num_freqs), ..., 2^(i / num_freqs), ...];
      Otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)].
  
      Args:
          num_freqs (int): the number of frequencies, default is 6;
          logspace (bool): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...],
              otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1)];
          input_dim (int): the input dimension, default is 3;
          include_input (bool): include the input tensor or not, default is True.
  
      Attributes:
          frequencies (torch.Tensor): If logspace is True, then the frequency f_i is [..., 2^(i / num_freqs), ...],
                  otherwise, the frequencies are linearly spaced between [1.0, 2^(num_freqs - 1);
  
          out_dim (int): the embedding size, if include_input is True, it is input_dim * (num_freqs * 2 + 1),
              otherwise, it is input_dim * num_freqs * 2.
  
      """
  
      def __init__(
          self,
          num_freqs: int = 6,
          logspace: bool = True,
          input_dim: int = 3,
          include_input: bool = True,
          include_pi: bool = True,
      ) -> None:
          """The initialization"""
  
          super().__init__()
  
          if logspace:
              frequencies = 2.0 ** torch.arange(num_freqs, dtype=torch.float32)
          else:
              frequencies = torch.linspace(
                  1.0, 2.0 ** (num_freqs - 1), num_freqs, dtype=torch.float32
              )
  
          if include_pi:
              frequencies *= torch.pi
  
          self.register_buffer("frequencies", frequencies, persistent=False)
          self.include_input = include_input
          self.num_freqs = num_freqs
  
          self.out_dim = self.get_dims(input_dim)
  
      def get_dims(self, input_dim):
          temp = 1 if self.include_input or self.num_freqs == 0 else 0
          out_dim = input_dim * (self.num_freqs * 2 + temp)
  
          return out_dim
  
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          """Forward process.
  
          Args:
              x: tensor of shape [..., dim]
  
          Returns:
              embedding: an embedding of `x` of shape [..., dim * (num_freqs * 2 + temp)]
                  where temp is 1 if include_input is True and 0 otherwise.
          """
  
          if self.num_freqs > 0:
              embed = (x[..., None].contiguous() * self.frequencies).view(
                  *x.shape[:-1], -1
              )
              if self.include_input:
                  return torch.cat((x, embed.sin(), embed.cos()), dim=-1)
              else:
                  return torch.cat((embed.sin(), embed.cos()), dim=-1)
          else:
              return x
  