Векторизованное создание массива диагональных квадратных массивов из массива лайнеров в Numpy или Tensorflow

Question

У меня есть массив формы [batch_size, N], например:

[[1  2]
 [3  4]
 [5  6]]

, и мне нужно создать массив из 3 индексов с формой [batch_size, N, N], где для каждого batch у меня есть N x N диагональная матрица, где диагонали взяты соответствующим элементом batch, например, в этом случае, в этом простом случае, результат, который я ищу:

[
  [[1,0],[0,2]],
  [[3,0],[0,4]],
  [[5,0],[0,6]],
]

Как я могу сделать эту операциюбез петель и взрывной векторизации?Я предполагаю, что это расширение измерения, но я не могу найти правильную функцию для этого.(Мне это нужно, так как я работаю с tenorflow и создаю прототипы с numpy).

Answer 1

Подход № 1

Вот векторизованный вариант с np.einsum для входного массива, a -

# Initialize o/p array
out = np.zeros(a.shape + (a.shape[1],),dtype=a.dtype)

# Get diagonal view and assign into it input array values
diag = np.einsum('ijj->ij',out)
diag[:] = a

Подход № 2

Другой, основанный на нарезке для назначения -

m,n = a.shape
out = np.zeros((m,n,n),dtype=a.dtype)
out.reshape(-1,n**2)[...,::n+1] = a

Answer 2

Использование np.expand_dims для поэлементного произведения с np.eye

a = np.array([[1,  2],
              [3,  4],
              [5, 6]])
N = a.shape[1]
a = np.expand_dims(a, axis=1)
a*np.eye(N)

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

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

      [[5., 0.],
       [0., 6.]]])

Пояснение

np.expand_dims(a, axis=1) добавляетновая ось в a, которая теперь будет (3, 1, 2) ndarray:

array([[[1, 2]],

       [[3, 4]],

       [[5, 6]]])

Теперь вы можете умножить этот массив на размерную матрицу N, которую вы можете сгенерировать с помощью np.eye:

np.eye(N)
array([[1., 0.],
       [0., 1.]])

, что даст желаемый результат:

a*np.eye(N)

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

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

       [[5., 0.],
        [0., 6.]]])

Answer 3

Попробуйте в тензорном потоке:

import tensorflow as tf
A = [[1,2],[3 ,4],[5,6]]
B = tf.matrix_diag(A)
print(B.eval(session=tf.Session()))
[[[1 0]
  [0 2]]

 [[3 0]
  [0 4]]

 [[5 0]
  [0 6]]]

Answer 4

Юй может использовать numpy.diag

m = [[1, 2],
 [3, 4],
 [5, 6]]

[np.diag(b) for b in m]

---

РЕДАКТИРОВАТЬ На следующем графике показано среднее время выполнения для решения выше (сплошная линия)и сравнил его с @ Divakar's (пунктирная линия) для разных размеров партии и разных размеров матрицы

Не думаю, что вы получите многоулучшение, но это только на основе этой простой метрики