What does axis = 0 do in Numpy's sum function?

If someone need this visual description:

numpy axis 0 and axis 1


All that is going on is that numpy is summing across the first (0th) and only axis. Consider the following:

In [2]: a = np.array([1, 2, 3])

In [3]: a.shape
Out[3]: (3,)

In [4]: len(a.shape) # number of dimensions
Out[4]: 1

In [5]: a1 = a.reshape(3,1)

In [6]: a2 = a.reshape(1,3)

In [7]: a1
Out[7]: 
array([[1],
       [2],
       [3]])

In [8]: a2
Out[8]: array([[1, 2, 3]])

In [9]: a1.sum(axis=1)
Out[9]: array([1, 2, 3])

In [10]: a1.sum(axis=0)
Out[10]: array([6])

In [11]: a2.sum(axis=1)
Out[11]: array([6])

In [12]: a2.sum(axis=0)
Out[12]: array([1, 2, 3])

So, to be more explicit:

In [15]: a1.shape
Out[15]: (3, 1)

a1 is 2-dimensional, the "long" axis being the first.

In [16]: a1[:,0] # give me everything in the first axis, and the first part of the second
Out[16]: array([1, 2, 3])

Now, sum along the first axis:

In [17]: a1.sum(axis=0)
Out[17]: array([6])

Now, consider a less trivial two-dimensional case:

In [20]: b = np.array([[1,2,3],[4,5,6]])

In [21]: b
Out[21]: 
array([[1, 2, 3],
       [4, 5, 6]])

In [22]: b.shape
Out[22]: (2, 3)

The first axis is the "rows". Sum along the rows:

In [23]: b.sum(axis=0)
Out[23]: array([5, 7, 9])

The second axis are the "columns". Sum along the columns:

In [24]: b.sum(axis=1)
Out[24]: array([ 6, 15])

The axis i in np.sum(a, axis=i) is the ith index of the shape of that array (zero-indexed).

Let's try and understand what that means with some examples:

a = np.array([1, 2, 3])
print (a.shape) #prints (3,) 
#so axis = 0 corresponds to 3 and axis = 1 corresponds to nothing

Let's see what axis = 0 and axis = 1 do to the sum:

sum = np.sum(a, axis=0) #sum = 6

So, sum = np.sum(a, axis=0) would sum up all numbers that the 0th index of a.shape refers to, which in this case are 3 numbers. Since numpy arrays are, by default, row-major (which is just another way of saying that the row index is specified before the column index), thus axis=0 would sum the three numbers that the shape refers to.

sum = np.sum(a, axis=1) #gives an error

Similarly, np.sum(a, axis=1) should sum up all the numbers that the 1st index of np.shape refers to, but since there is no first index of the shape, we get an error.

Let's take another example:

b = np.array([[1,2,3],
             [4,5,6]])
print(b.shape) #prints (2,3)
#axis = 0 corresponds to 2 and axis = 1 corresponds to 3

And now, let's see what changing the axis does:

sum = np.sum(b, axis=0) #sum = [5, 7, 9] of shape(3,)

We know that axis = 0 should sum along the first index of the shape and we expect it to find two numbers along this axis (by looking at the shape). So [1+4, 2+5, 3+6].

sum = np.sum(b, axis=1) #sum = [6, 15] of shape(2,)

Now the sum is along axis = 1, and from the shape we can see this it is an axis along which there are 3 numbers to be summed. So, [1+2+3,4+5+6]