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.

Thanks in advanced!