python - arrayfire-python の行/列方向の外積

Question

arrayfire-pythonで行/列ごとの外積を実装するには?

これはnumpyでの類推実装です。

A = np.random.randint(1,7, size=(3,2))
B = np.random.randint(1,7, size=(2,3))
X = np.einsum('ik,kj->kij', A,B)
X
array([[[ 8, 16, 12],
        [ 8, 16, 12],
        [ 4,  8,  6]],

       [[15,  3, 12],
        [ 5,  1,  4],
        [ 5,  1,  4]]])

PS ミニバッチ勾配降下法を実装するためのわずかなスピードアップを探していて、このライブラリを見つけました。