Positional Encoding Layer that allows for Batched Inputs

saethrej: <p>Hi everyone. I have written a custom module for a Positional Encoding layer in a neural network. However, this one only works for “flat” inputs x, such as individual indices, rather than e.g. entire tensors with multiple dimensions.</p>
<pre><code class="lang-auto">class PositionalEncoding(torch.nn.Module):

    def __init__(self, b: float = 1.25, l: int = 80) -&gt; None:
        super(PositionalEncoding, self).__init__()
        self.l = int(l)
        self.b = b

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        x_pe = []
        for i in range(self.l):
            arg = (self.b ** i) * torch.pi * x
            x_pe.extend([torch.sin(arg), torch.cos(arg)])
        return torch.stack(x_pe, dim = 1).squeeze(-1)
</code></pre>
<p>My goal now is to obtain a module that can positionally encode entire batches of tensors at once. An example would e.g. be a <code>[B, N, K]</code> tensor, where B is the batch size. The desired output shape is <code>[B, N, 2*K*self.l]</code>. I have tried this for a few hours, and I seem to be no closer to a solution that is functionally equivalent to the above.</p>
<p><strong>My main idea:</strong> I would store the input shape of <code>x</code>, and then flatten <code>x</code> except for the last dimension, and unsqueeze a new dimension at the second-to-last dimension for frequncies. I<code>d then compute all the frequencies as a tensor with equally many dimensions as </code>x<code>, and multiple </code>x<code>with the frequencies to get a tensor that is conceptually the same as</code>arg<code>above. I then want to take the sine and the cosine of this tensor, and stack them in some way such that for every element of the tensor that is positionally encoded, the order of the elements in its embeddings are</code>[sin(x1), cos(x1), sin(x2), cos(x2), …]`.</p>
<p><strong>Problem:</strong> I can’t find a way that achieves this interleaving of the sines and the cosines… I tried with <code>stack()</code>, <code>view()</code>, <code>reshape()</code>, <code>permute()</code>, <code>dstack()</code> and multiple variations and combinations thereof, but to no avail. I also tried using <code>einops.rearrange</code> but also did not get the desired result.</p>
<p>For reference, this is an approach I tried which returns something similar, but not quite correct:</p>
<pre><code class="lang-auto">def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        # store input shape to later recover it and flatten all dimensions, except the last one,
        # and add a new dimensions for the frequencies
        old_shape = x.shape
        x_flat = x.view(-1, 1, old_shape[-1])

        # compute the exponentially increasing frequences and multiply them with the flattened input
        # using PyTorch's broadcasting functionality
        freqs = torch.pi * (self.b ** torch.arange(self.l, dtype = torch.float32, device = x.device))
        arg = x_flat * freqs.view(1, self.l, 1)

        # take the sine and the consine of the arguments, and interleave them in the last dimension
        # we want [sin(x1), cos(x1), sin(x2), cos(x2), ...] , not [sin(x1), sin(x2), ..., cos(x1), cos(x2), ...]
        # to do so, we stack the sines and the cosines in the last dimension and then flatten the
        # last two dimensions.
        x_pe_int_extradim = torch.dstack([torch.sin(arg), torch.cos(arg)])
        x_pe_int = x_pe_int_extradim.view(x_pe_int_extradim.size(0), -1)
        x_pe = x_pe_int.view(*old_shape[:-1], -1)
        return x_pe
</code></pre>
<p>I’d very much appreciate any help or pointers to resources that can help me!</p>

Jonathan_Coletti: <p>To perform this you can use the class and function below (taken from <a href="https://github.com/tatp22/multidim-positional-encoding" class="inline-onebox" rel="noopener nofollow ugc">GitHub - tatp22/multidim-positional-encoding: An implementation of 1D, 2D, and 3D positional encoding in Pytorch and TensorFlow</a>) (so all credit to him) within the repo there is also 1D, 2D, 3D, along with permuted versions. Anyway here is the snippet</p>
<pre><code class="lang-auto">def get_emb(sin_inp):
    """
    Gets a base embedding for one dimension with sin and cos intertwined
    """
    emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
    return torch.flatten(emb, -2, -1)
</code></pre>
<pre><code class="lang-auto">class PositionalEncoding1D(nn.Module):
    def __init__(self, channels, dtype_override=None):
        """
        :param channels: The last dimension of the tensor you want to apply pos emb to.
        :param dtype_override: If set, overrides the dtype of the output embedding.
        """
        super(PositionalEncoding1D, self).__init__()
        self.org_channels = channels
        channels = int(np.ceil(channels / 2) * 2)
        inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("cached_penc", None, persistent=False)
        self.channels = channels
        self.dtype_override = dtype_override

    def forward(self, tensor):
        """
        :param tensor: A 3d tensor of size (batch_size, x, ch)
        :return: Positional Encoding Matrix of size (batch_size, x, ch)
        """
        if len(tensor.shape) != 3:
            raise RuntimeError("The input tensor has to be 3d!")

        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
            return self.cached_penc

        self.cached_penc = None
        batch_size, x, orig_ch = tensor.shape
        pos_x = torch.arange(x, device=tensor.device, dtype=self.inv_freq.dtype)
        sin_inp_x = torch.einsum("i,j-&gt;ij", pos_x, self.inv_freq)
        emb_x = get_emb(sin_inp_x)
        emb = torch.zeros(
            (x, self.channels),
            device=tensor.device,
            dtype=(
                self.dtype_override if self.dtype_override is not None else tensor.dtype
            ),
        )
        emb[:, : self.channels] = emb_x

        self.cached_penc = emb[None, :, :orig_ch].repeat(batch_size, 1, 1)
        return self.cached_penc
</code></pre>
<p>It uses caching along with manual repeating along the batch size. Furthermore, there exists a different frequency scaling (see <code>inv_freq</code>) so you may have to change that</p>