Hi, I am looking for a solution to achieve the effect of expanding a 3x4 tensor on its diagonal axis and fill the new diagonal with 1. The original tensor is like a weighting map with each column sum up to 1:
[[0.5,  0.8,  0.7,  0.3],
 [0.2,  0.1,  0.2,  0.4],
 [0.3,  0.1,  0.1,  0.3]]

and the target matrix (y) I want to expand to looks like:
[[1,    0.8,  0.7,  0.3],
 [0.5,  1,    0.2,  0.4],
 [0.2,  0.1,  1,    0.3],
 [0.3,  0.1,  0.1,  1   ]]

and the underlying reason that I want to achieve this is because I have an input tensor x with 4x4 dimension, and I would like to produce a 4x4 matrix by assigning the weightings obtained from y to the apply to the input tensor x (like the attention mechanism), so the resulting 4x4 is like this in pseudo:
[
 [ Y11 * X_row1 + Y21 * X_row2 + Y31 * X_row3 + Y41 * X_row4 ],
 [ Y12 * X_row1 + Y22 * X_row2 + Y32 * X_row3 + Y44 * X_row4 ]
 [ Y13 * X_row1 + Y23 * X_row2 + Y33 * X_row3 + Y43 * X_row4 ],
 [ Y14 * X_row1 + Y24 * X_row2 + Y34 * X_row3 + Y44 * X_row4 ]
]

which can be illustrated as:
Screenshot from 2021-01-20 15-16-311126×542 38.8 KB
I appreciate any thoughts which related to (1) how to expand the tensor on the diagonal axis, or thoughts regarding (2) a better way of achieving the resulting attention mechanism with an easier way which won’t requires the expansion of Y.
Thanks in advanced!