Selecting indices based on values of tensor and then use it for torch.gather

psakamoori: <p>I have two tensors<br>
X: torch.Size([16, 10, 100])<br>
Mask: torch.Size([16, 10, 20])</p>
<p>16 → Batch_size<br>
10 → total ID’s<br>
100 → Total events for each ID<br>
20 → Length of each event</p>
<p>Requirement: How to create 1D indices matrix - i.e., indices will have/store last index value of mask tensor at which its cell value = 1.   For ex: Mask[0][0] = [1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]  → last index at which value = 1 is indx=3. So, I need to have 1D array of indices as = [[3,…]]</p>
<p>Once I have 1D indices I need to select values in “X” at those indices using torch.gather.<br>
Appreciate any pointers…on how I can use this efficiently.</p>

KFrank: <p>Hi Pradeep!</p>
<aside class="quote no-group" data-username="psakamoori" data-post="1" data-topic="144948" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/psakamoori/48/46995_2.png" class="avatar"> psakamoori:</div>
<blockquote>
<p>Once I have 1D indices I need to select values in “X” at those indices using torch.gather.</p>
</blockquote>
</aside>
<p>Could you post a fully-self-contained, runnable script that performs an<br>
example of the calculation you want to do, using python <code>for</code> loops and<br>
<code>if</code> blocks, as necessary, together with its output?</p>
<p>Best.</p>
<p>K. Frank</p>

psakamoori: <p>For now, I just have problem statement - I am trying to build the script for the problem statement… But, looking to find right torch APIs for efficient implementation…Let me try explain…</p>
<p>Input: X 4D tensor</p>
<p>“X” is collection of health records of 10 patients - record per patient is a tensor of n x m (meaning each patient had “n” visits to the hospital and “m” is max number of diagnosis tests that performed per visi…So, x looks like below</p>
<p>X[batch_size][number_of_visits][max.no.of.diagnosis.performed.per.visit_for_a_patient_in_the_group_of_10][embedding_dimension]</p>
<p>Assuming below params:<br>
batch_size = 16<br>
no_of_visits = 10<br>
max.no.of.diagnosis.performed.per.visit_for_a_patient_in_the_group_of_10 = 20 (If any diagnosis code less than 20, that tensor will be padded with zeros  - will talk about this mask below).<br>
Embedding dimension = 100  (total records)</p>
<p>Mask: masks 3D tensor</p>
<p>masks[batch_size][number_of_visits][max.no.of.diagnosis.performed.per.visit_for_a_patient_in_the_group_of_10]</p>
<p>Sample values of X and Masks:<br>
X =  [ [ [10, 8, 21], [18, 5] ],<br>
[ [1, 3, 62, 17], [1], [3, 1] ] ]<br>
bs = 1<br>
no_of_visits = 2<br>
max_diagnosis_performed_in_a_visit = 4  (i.e, len([1, 3, 62, 17]))</p>
<pre><code class="lang-auto">
Padded "X" looks like below:
[ [ [10, 8, 21, *0*], [18, 5, *0*, *0*], [*0*, *0*, *0*, *0*]  ], 
  [ [1, 3, 62, 17], [1, *0*, *0*, *0*], [3, 1, *0*, *0*] ] ]

So size of padded X = [2][4][12]

Masks tensor for above X will be:
[ [ [1, 1, 1, 0], [1, 1, 0, 0], [0, 0, 0, 0] ], 
  [ [1, 1, 1, 1], [1, 0, 0, 0], [1, 1, 0, 0] ] ]

Now Question is: 
For above given padded "X" and "Masks" - I need to perform below

1. Extract all last true visit (that is index of last non-padded value) of each patient:

Ex: for Patient 1: it will be [[3, 2]] as per mask of patient:1 "[ [1, 1, 1, 0], [1, 1, 0, 0], [0, 0, 0, 0] ]" 
last true visit per visit is 1st row 3rd col, 2 row 2nd col and no true visit in 3rd row

2. For above indices extracted for each patient - I need to use this indices for all patients per-batch size (above example just talks with batch size =1, but I might be have 16/32 as bs) -  to construct a tensor from  "X"  using torch.gather()

Please let me know, if need more details.</code></pre>

Eta_C: <p>If values <code>1</code> in your <code>mask</code> is contiguous (there is not <code>0</code> between two <code>1</code>s) and all of each vector starts with <code>1</code>, I mean, only <code>[1, 1, 0, 0, 0]</code> is supported, but not <code>[0, 1, 1, 0, 0]</code> or <code>[1, 0, 1, 0, 0]</code>, you can try my following code.</p>
<pre><code class="lang-python">indices = torch.sum(mask, dim=-1, keepdim=True) - 1
</code></pre>

psakamoori: <p>Thank you for the code snippet Eta_C…I tried that- But, output size is not exactly what i am looking for - output size from code snipper is  indices = torch.Size([16, 10, 1]) …But, I am looking for indices with size [1, 16]</p>
<p>i.e.</p>
<ol>
<li>Input masks[batch_size][visits][total diagnosis]</li>
<li>Flatten it to mask[batch_size][visits * diagnosis]</li>
<li>Sum each row in mask[bastch_sizer] (i.e. all the (visits * diagnosis) values with each other per batch_size)</li>
<li>this will give me a final matrix/tensor of size mask[1, 16]</li>
</ol>
<p>Notes: Yes masks are contiguous</p>

psakamoori: <p>Able to implement what I am looking for.<br>
Not able to share the implementation here now…due to restrictions…But used below APIs</p>
<p>torch.sum(x, dim=2)<br>
torch.view<br>
torch.expand<br>
torch.gather<br>
torch.squeeze</p>