1

In python, assume

a = np.array(range(0,12)).reshape(2,2,3)
b = np.array(range(0,6)).reshape(3,2)
c = np.matmul(a,b) // a @ b

We have

a: array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])

b: array([[0, 1],
       [2, 3],
       [4, 5]])

c: array([[[10, 13],
        [28, 40]],

       [[46, 67],
        [64, 94]]])

Could someone help me to achieve equivalent operation in java nd4j without a for loop? I've tried broadcast.mul, but it turns out broadcast.mul is element-wise multiplication. I didn't found any broadcast operation for mmul.

Xiang Liu
  • 43
  • 6

1 Answers1

2

I figured it out by myself. The answer is shown below in case someone needs it. With Nd4j.tensorMmul, matrix broadcast could be easily achieved. e.g.

val a = Nd4j.create(0d to 11d by 1d toArray, Array[Int](2, 2, 3))
val b = Nd4j.create(0d to 5d by 1d toArray, Array[Int](3, 2))
Nd4j.tensorMmul(a, b, Array(Array(2), Array(0))) // matrix broadcast

It's code for scala. For java, you just need change the code to create arrays.

Xiang Liu
  • 43
  • 6