import numpy as np
import numba
@numba.vectorize('i4(i4)', target = 'parallel')
def mag(b):
return b * b
def main():
mat_a = np.full((5, 3),2,dtype=np.int32)
c = mag(mat_a)
d = np.sum(c, axis = 1)
print d
OUTPUT: [12 12 12 12 12]
But I want the output to be like this:
[12 12 12]
[12 12 12]
[12 12 12]
[12 12 12]
[12 12 12]
Clear example:
Suppose I have a output like this: [12 13 14 15 16] I want to convert each element in the vector into my own dimensions with numpy
[[12 12 12]
[13 13 13]
[14 14 14]
[15 15 15]
[16 16 16]]