row/column-wise outer product in arrayfire-python

Question

How to implement row/column-wise outer product in arrayfire-python?

Here is analogical implementation in 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]]])

P.S. I'm looking for slight speed up for implementation of Mini-batch Gradient Descent, and I've found this library.

Answer 1

This isn't available in arrayfire as a native function. You can however implement something using broadcast and moddims like this:

>>> @af.broadcast
... def outer(a, b):
...     am = af.moddims(a, a.shape[0], 1, a.shape[1])
...     bm = af.moddims(b.T, 1, b.shape[1], b.shape[0])
...     cm = am * bm
...     return cm
... 
>>> a
arrayfire.Array()
Type: float
[3 2 1 1]
    0.6874     0.3552 
    0.9058     0.1589 
    0.5023     0.8857 


>>> b
arrayfire.Array()
Type: float
[2 3 1 1]
    0.2060     0.9028     0.7167 
    0.4522     0.0121     0.2723 


>>> print(outer(a, b))
arrayfire.Array()
Type: float
[3 3 2 1]
    0.1416     0.6206     0.4927 
    0.1866     0.8178     0.6492 
    0.1035     0.4535     0.3600 

    0.1606     0.0043     0.0967 
    0.0719     0.0019     0.0433 
    0.4006     0.0107     0.2412 

If you have any particular feature requests, follow up on github: https://github.com/arrayfire/arrayfire-python/issues