[Attention Mechanism] Is there a way of expanding a tensor on its diagonal axis?

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:

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.