Normalise list of N dimensional numpy arrays

Question

I have a list of N dimensional NumPy arrays.

num_vecs = 10
dims = 2

vecs = np.random.normal(size=(num_vecs, dims))

I want to normalize them, so the magnitude/length of each vector is 1. I can easily do this with a for-loop.

norm = np.linalg.norm(vecs, axis=1)

for dd in range(dims):
    vecs[:, dd] /= norm

assert np.allclose(np.linalg.norm(vecs, axis=1), 1.)

But how do I get rid of the for-loop?

Answer 1

The trick is to use the keepdims parameter.

import numpy as np

num_vecs = 10
dims = 2

vecs = np.random.normal(size=(num_vecs, dims))


vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)

assert np.allclose(np.linalg.norm(vecs, axis=1), 1.)