Vectorized softmax calculation using numpy

softmax.py

import numpy as np

def softmax(x):
  """ Compute the softmax for each row of the input x
  Arguments:
  x -- A N dimensional veector or M X N dimensional numpy matrix.
  Return:
  x -- modified x in-place
  """
  
  if len(x.shape) > 1:
    #Matrix
    max_element = np.max(x, axis = 1)
    x -= max_element[:, None] #Broadcasting
    x = np.exp(x)
    x = x / np.sum(x, axis = 1)[:, None]
  else:
    #Vector
    max_element = np.max(x)
    x -= max_element
    x = np.exp(x)
    x = x / np.sum(x)
  return x