Indexing by trainable weights

I_H_Yoo: <p>The picture below describes calculation of the local attention of Luong Attention.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://img1.daumcdn.net/thumb/R1280x0/?scode=mtistory2&amp;fname=https%3A%2F%2Fk.kakaocdn.net%2Fdn%2Fd2leeT%2FbtquXku1VUH%2FcV90Ye7X1657nzphcOjQYk%2Fimg.png" title="" rel="nofollow noopener"><img src="https://img1.daumcdn.net/thumb/R1280x0/?scode=mtistory2&amp;fname=https%3A%2F%2Fk.kakaocdn.net%2Fdn%2Fd2leeT%2FbtquXku1VUH%2FcV90Ye7X1657nzphcOjQYk%2Fimg.png" alt="" width="690" height="388"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use xlink:href="#far-image"></use></svg><span class="filename"></span><span class="informations">770×433</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use xlink:href="#discourse-expand"></use></svg></div></a></div></p>
<p>What I’m trying to do is getting the alignment vector in the picture. In the initializer of my decoder, I define v_p^T and W_p which are trainable weights like this:</p>
<pre><code class="lang-python">self.attn_dim = hid_dim // 8
self.local_v = nn.Parameter(torch.randn(self.attn_dim, 1))
self.local_w = nn.Linear(hid_dim, self.attn_dim)
</code></pre>
<p>And then, I get alignment vector \alpha_{t, i} but the vector including all of range of the input sequence. Therefore, I need to sample it by (Pt + D, Pt - D) as showed in the picture.</p>
<pre><code class="lang-python">align_score  = self.attention(H, s) # [batch_size, 1, seq_len] # it should be [Batch_size, 1, 2D+1]
</code></pre>
<p>For the sampling, I need to get Pt. So I feed to the weights I defined in the initializer, and got Pt.</p>
<pre><code class="lang-python">temp = self.local_w(s.unsqueeze(1)) # [Batch, 1, attn_dim]
Pt = Tx * torch.sigmoid(temp.bmm(self.local_v.repeat(batch_size, 1, 1))) # [Batch, 1, 1]
</code></pre>
<p><code>Pt</code> is a tensor. I belive the elements in <code>Pt</code> are different each other, in every batch iteration. So the vector must be form of a tensor. However, at the same time, I also get <code>align_score</code> within the window [Pt-D, Pt+D] by using Pt tensor.</p>
<p>What I’m struggling is that the both conditions which is the grad must be flew into Pt (because of v_p and W_p), and which is the another condition that the Pt tensor should be utilized to getting index of the window within [Pt-D, Pt+D] should be satisfied at the same time. Can someone give me an advice? The entire code of calculating alignment score is shown below.</p>
<pre><code class="lang-python">temp = self.local_w(s.unsqueeze(1)) # [Batch, 1, attn_dim]
Pt = Tx * torch.sigmoid(temp.bmm(self.local_v.repeat(batch_size, 1, 1))) # [Batch, 1, 1]

start = ??
end =  ??
align_score = align_score[:, :, start:end] # [batch_size, 1, 2D + 1]
local = torch.tensor([x for x in range(1, align_score.size(2) + 1)]) # [2D + 1]
local = torch.exp(- ((local -   Pt) ** 2)/((self.D ** 2) / 2)).view(1, 1, -1) # [2D + 1]
local = local.repeat(batch_size, 1, 1) # [B, 1, 2D + 1]
align_score = align_score * local  # [B, 1, 2D + 1]
context = torch.bmm(align_score, H.permute(1, 0, 2)[:, start:end, :]) # [batch_size, 1, enc_hid]
</code></pre>