Create binary tensor from index tensor

sprakashdash: <p>I have an index tensor like we send input to a bert model after tokenization, so the shape is <code>[batch_size, seq_len]</code> and I have to return a tensor with <code>0</code>’s and <code>1</code>’s (<code>1</code>’s where the indices are matched) of shape <code>[batch_size, seq_len, vocab_size]</code>.</p>

sprakashdash: <p>The closest solution I could find is to use <a href="https://pytorch.org/functorch/stable/generated/functorch.vmap.html" rel="noopener nofollow ugc">vmap</a>.<br>
I defined a function to take an index and dimension and output a <code>torch.zeros</code> of the given dim (which is a 1-dimensional for my case) and add 1 to a specific index.</p>
<pre><code class="lang-auto">def idx2tensor(val, dim):         # apply this across second dimension
            empty_tensor = torch.zeros(dim)
            empty_tensor[val] = 1
</code></pre>
<p>Now I am applying vmap for an input tensor of shape <code>[batch_size, seq_length]</code> and expect to get <code>[batch_size, seq_length, dim]</code>.<br>
<code>batched_idx2ten = torch.vmap(idx2tensor, in_dims=(0, None))</code> following <a href="https://pytorch.org/functorch/stable/generated/functorch.vmap.html#:~:text=If%20there%20are%20multiple%20inputs%20each%20of%20which%20is%20batched%20along%20different%20dimensions%2C%20in_dims%20must%20be%20a%20tuple%20with%20the%20batch%20dimension%20for%20each%20input%20as" rel="noopener nofollow ugc">this example in docs</a>.<br>
Apply the <code>batched_idx2ten</code> to a random input and a fixed dim:</p>
<pre><code class="lang-auto">x = 10 * torch.rand(4, 10) - 1
x = torch.tensor(x, dtype=torch.int64)
batched_idx2ten(x, 15)
</code></pre>
<p>I am now stuck with the following error:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-26-80c94fc2eaa3&gt; in &lt;cell line: 2&gt;()
      1 x = 10 * torch.rand(4, 10) - 1
----&gt; 2 batched_idx2ten(x, 15)

4 frames
/usr/local/lib/python3.10/dist-packages/torch/_functorch/apis.py in wrapped(*args, **kwargs)
    186     # @functools.wraps(func)
    187     def wrapped(*args, **kwargs):
--&gt; 188         return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
    189 
    190     return wrapped

/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
    276 
    277     # If chunk_size is not specified.
--&gt; 278     return _flat_vmap(
    279         func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs
    280     )

/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)
     42     def fn(*args, **kwargs):
     43         with torch.autograd.graph.disable_saved_tensors_hooks(message):
---&gt; 44             return f(*args, **kwargs)
     45     return fn
     46 

/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)
    389     try:
    390         batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
--&gt; 391         batched_outputs = func(*batched_inputs, **kwargs)
    392         return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)
    393     finally:

&lt;ipython-input-25-c7783e629350&gt; in idx2tensor(val, dim)
      1 def idx2tensor(val, dim):         # apply this across second dimension
      2     empty_tensor = torch.zeros(dim)
----&gt; 3     empty_tensor[val] = 1
      4 
      5 batched_idx2ten = torch.vmap(idx2tensor, in_dims=(0, None))

RuntimeError: vmap: index_put_(self, *extra_args) is not possible because there exists a Tensor `other`
in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but
`self` not being vmapped over in a vmap. Please try to use out-of-place operators instead of index_put_. 
If said operator is being called inside the PyTorch framework, please file a bug report instead.
</code></pre>

KFrank: <p>Hi Satya!</p>
<aside class="quote no-group" data-username="sprakashdash" data-post="2" data-topic="200269" 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/sprakashdash/48/48780_2.png" class="avatar"> sprakashdash:</div>
<blockquote>
<pre><code class="lang-auto">def idx2tensor(val, dim):         # apply this across second dimension
            empty_tensor = torch.zeros(dim)
            empty_tensor[val] = 1
</code></pre>
</blockquote>
</aside>
<p>I believe that you are looking for <code>torch.nn.functional.one_hot()</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.2.2'
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt; x = 10 * torch.rand(4, 10) - 1
&gt;&gt;&gt; x = torch.tensor(x, dtype=torch.int64)
&lt;stdin&gt;:1: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
&gt;&gt;&gt; x
tensor([[4, 7, 8, 0, 0, 2, 1, 5, 3, 6],
        [0, 4, 0, 6, 8, 3, 2, 8, 6, 5],
        [7, 8, 5, 5, 1, 4, 1, 6, 0, 0],
        [1, 0, 8, 5, 5, 3, 4, 0, 2, 4]])
&gt;&gt;&gt; result = torch.nn.functional.one_hot (x, 15)
&gt;&gt;&gt; result.shape
torch.Size([4, 10, 15])
&gt;&gt;&gt; result[0]
tensor([[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>