Create a 3D (partially) diagonal array from a 2D array-CodePudding

I'd like to ask how can I efficiently generate a numpy 3D array from a 2D array with each row filling the diagonal part of the new array? For example, the input 2D array is

array([[1, 2],
       [3, 4],
       [5, 6],
       [7, 8]])

and I want the output to be

array([[[1, 0],
        [0, 2]],

       [[3, 0],
        [0, 4]],

       [[5, 0],
        [0, 6]],

       [[7, 0],
        [0, 8]]])

Typically, the size of the first dimensional is very large. Thanks in advance.

CodePudding user response：

Assuming a the input and using indexing with

CodePudding user response：

arr = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
res = np.apply_along_axis(np.diag, 1, arr)