How to create a 3d indexable embedding matrix?

peternasser99: <p>import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F</p>
<p>def generate_and_add_binary_numbers(num_samples, bit_size, num2_decimal=None):<br>
num1_bin = [‘’.join(str(bit) for bit in torch.randint(0, 2, (bit_size,)).tolist()) for _ in range(num_samples)]<br>
if num2_decimal is None:<br>
num2_bin = [‘’.join(str(bit) for bit in torch.randint(0, 2, (bit_size,)).tolist()) for _ in range(num_samples)]<br>
else:<br>
# Convert the decimal number to a binary string with padding<br>
num2_bin_str = format(num2_decimal, ‘b’).zfill(bit_size)[-bit_size:]<br>
num2_bin = [num2_bin_str] * num_samples  # Repeat for all samples</p>
<pre><code># Convert binary strings to integers for addition
int1 = [int(n, 2) for n in num1_bin]
int2 = [int(n, 2) for n in num2_bin]
sum_int = [i1 + i2 for i1, i2 in zip(int1, int2)]

# Convert the sums back to binary strings with padding
sum_bin = [format(s, 'b').zfill(bit_size)[-bit_size:] for s in sum_int]

return num1_bin, num2_bin, sum_bin
</code></pre>
<p>def bin_list_to_tensor(*bin_lists):<br>
tensor_lists = [<span class="chcklst-box fa fa-square-o fa-fw"></span> for _ in range(len(bin_lists))]<br>
for i, bin_list in enumerate(bin_lists):<br>
for binary_str in bin_list:<br>
binary_list = [int(bit) for bit in binary_str]<br>
binary_tensor = torch.tensor(binary_list, dtype=torch.float32)<br>
tensor_lists[i].append(binary_tensor)<br>
tensor_lists[i] = torch.stack(tensor_lists[i])  # Convert list of tensors to a tensor<br>
return tensor_lists</p>
<p>num_samples = 1000<br>
bit_size = 8<br>
num_epochs = 5000<br>
learning_rate = 0.001</p>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F</p>
<p>class AttentionModule(nn.Module):<br>
def <strong>init</strong>(self, n_embd , bit_size = bit_size):<br>
super(AttentionModule, self).<strong>init</strong>()</p>
<pre><code>    self.key = nn.Linear(n_embd, n_embd, bias=False)
    self.query = nn.Linear(n_embd, n_embd, bias=False)
    self.value = nn.Linear(n_embd, n_embd, bias=False)
    
    self.ln1 = nn.LayerNorm(n_embd)

def forward(self, x):
    
    k = self.key(x)
    q = self.query(x)
    v = self.value(x)
    
    wei = q @ k.transpose(-2,-1) * n_embd**-0.5 
    wei = F.softmax(wei, dim=-1)
    out = wei @ v
    
    out = self.ln1(out + x)
    
    return out
</code></pre>
<p>class block(nn.Module):<br>
def <strong>init</strong>(self, n_embd , bit_size = bit_size , num_layers = 3 ):<br>
super(block, self).<strong>init</strong>()<br>
self.embedding_layer = nn.Parameter(torch.randn(size=(bit_size, 2, n_embd)))<br>
self.embedding_layer2 = nn.Parameter(torch.randn(size=(bit_size, 2, n_embd)))<br>
self.att_blocks = nn.ModuleList([<br>
nn.Sequential(<br>
AttentionModule(n_embd),</p>
<h1><a name="nnrelu-1" class="anchor" href="#nnrelu-1"></a>nn.ReLU()</h1>
<pre><code>            nn.Sigmoid()
        )
        for _ in range(num_layers)
    ])
    
    self.fc = nn.Linear(n_embd * bit_size *2, bit_size, bias=False)

    
def forward(self, num1 , num2):
    tokenized_num1 = num1.long()
    tokenized_num2 = num2.long()
    embedded_num1 = self.embedding_layer[torch.arange(num1.shape[0]).unsqueeze(1), tokenized_num1]
    embedded_num2 = self.embedding_layer2[torch.arange(num2.shape[0]).unsqueeze(1), tokenized_num2]
    
    print(embedded_num1.shape)
    print()
    
    x = torch.cat((embedded_num1, embedded_num2), dim=1)
    
    for b in self.att_blocks:
        x = b(x)
        
    x = x.view(x.size(0), -1)
    x = self.fc(x)
    return torch.sigmoid(x)
</code></pre>
<p>n_embd = 4<br>
model = block(n_embd).to(device)</p>
<p>num1, num2, sum_bin = generate_and_add_binary_numbers(num_samples, bit_size , num2_decimal=2)<br>
num1, num2, sum_bin = bin_list_to_tensor(num1, num2, sum_bin)<br>
num1, num2, sum_bin = num1.to(device), num2.to(device), sum_bin.to(device)</p>
<p>out = model(num1, num2)</p>
<p>print(“out:”, out)</p>
<p>those 2 lines give cuda runtime error<br>
one runs normally but rerunning it or running the second breaks it<br>
embedded_num1 = model.embedding_layer[torch.arange(num1.shape[0]).unsqueeze(1), num1.long()]<br>
embedded_num2 = model.embedding_layer2[torch.arange(num2.shape[0]).unsqueeze(1), num2.long()]</p>