Тензорный эквивалент этого числового осевого декартова произведения для двумерных матриц

Question

В настоящее время у меня есть код, который позволяет перемещать комбинаторное (декартово) произведение по определенной оси. Это в numpy, и возник из предыдущего вопроса Эффективное по оси координат декартово произведение нескольких 2D матриц с Numpy или TensorFlow

A = np.array([[1,2],
              [3,4]])
B = np.array([[10,20],
              [5,6]])
C = np.array([[50, 0],
              [60, 8]])
cartesian_product( [A,B,C], axis=1 )
>> np.array([[ 1*10*50, 1*10*0, 1*20*50, 1*20*0, 2*10*50, 2*10*0, 2*20*50, 2*20*0] 
             [ 3*5*60,  3*5*8,  3*6*60,  3*6*8,  4*5*60,  4*5*8,  4*6*60,  4*6*8]])

и повторить решение:

L = [A,B,C]  # list of arrays
n = L[0].shape[0]
out = (L[1][:,None]*L[0][:,:,None]).reshape(n,-1)
for i in L[2:]:
    out = (i[:,None]*out[:,:,None]).reshape(n,-1)

Существует ли существующий способ выполнить это с трансляцией в тензорном потоке - без цикла for?

Answer 1

Итак, мне удалось найти чистый (основанный на tf) ответ для двух массивов. В настоящее время оно не обобщается, как простое решение для массивов M, но это уже другой вопрос (возможно, tf.while_loop). Для тех, кто любопытен, решение адаптируется из оценки всех парных комбинаций строк двух тензоров в тензорном потоке

a = np.array([[0, 1, 2, 3],
              [4, 5, 6, 7],
              [4, 5, 6, 7]])

b = np.array([[0, 1],
              [2, 3],
              [2, 3]])

N = a.shape[0]

A = tf.constant(a, dtype=tf.float64)
B = tf.constant(b, dtype=tf.float64)

A_ = tf.expand_dims(A, axis=1)
B_ = tf.expand_dims(B, axis=2)
z = tf.reshape(tf.multiply(A_, B_), [N, -1])

>> tf_result
Out[1]: 
array([[  0.,   0.,   0.,   0.,   0.,   1.,   2.,   3.],
       [  8.,  10.,  12.,  14.,  12.,  15.,  18.,  21.],
       [  8.,  10.,  12.,  14.,  12.,  15.,  18.,  21.]])

Решения для случая с несколькими массивами приветствуются