split_linear.py

import torch
import torch.nn.functional as F

d = 8
seq_len = 13
bs = 1
wt = torch.rand((d, d))
x = torch.rand((seq_len, bs, d))
x_r0, x_r1 = x[:,:, :d//2], x[:,:, d//2:]
wt_r0, wt_r1 = wt[:, :d//2], wt[:, d//2:]

x_catted = torch.cat([x_r0, x_r1], 2)
wt_catted = torch.cat([wt_r0, wt_r1], 1)
result_catted = F.linear(x_catted, wt_catted)
#
or0_proj = F.linear(x_r0, wt_r0)
or1_proj = F.linear(x_r1, wt_r1)
result_summed = or0_proj + or1_proj
equality_pct = result_summed.eq(result_catted).float().mean()  # .6063
print(f'delta: {(result_summed - result_catted).abs().mean()}, equality_pct: {equality_pct}')
# delta: 5.788528056882569e-08, equality_pct: 0.6346153616905212