python 3 how to use multithreading with numpy linalg norm

Question

I am using python3 with np.linalg.norm to calculate the norms for rows in a matrix (norm(axis=1)), Is there a straightforward way, using only np to make it run using multithreading or multicoring?

Answer 1

We can use numexpr module that supports multi-core processing, like so -

import numexpr as ne

def linalg_norm(a):
    sq_norm = ne.evaluate('sum(a**2,1)')
    return ne.evaluate('sqrt(sq_norm)')

To perform norm-reduction along any other axis, replace 1 with that axis number in the evaluate expression - 'sum(a**2,1)'.

Sample run -

In [34]: np.random.seed(0)
    ...: a = np.random.rand(4,5)

In [35]: np.linalg.norm(a,axis=1)
Out[35]: array([1.28545589, 1.57467272, 1.4460319 , 1.43656019])

In [36]: linalg_norm(a)
Out[36]: array([1.28545589, 1.57467272, 1.4460319 , 1.43656019])

Related post on how to control multi-core functionality.

---

For completeness, few alternatives could be proposed.

An efficient solution would be with np.einsum -

In [39]: np.sqrt(np.einsum('ij,ij->i',a, a))
Out[39]: array([1.28545589, 1.57467272, 1.4460319 , 1.43656019])

With np.matmul/@ operator on Python 3.x -

In [6]: np.sqrt(np.matmul(a[:,None],a[:,:,None])[:,0,0])
Out[6]: array([1.28545589, 1.57467272, 1.4460319 , 1.43656019])

In [7]: np.sqrt((a[:,None] @ a[:,:,None])[:,0,0])
Out[7]: array([1.28545589, 1.57467272, 1.4460319 , 1.43656019])