title: How to vectorize Softmax probability of a multi dimensional matrix
body: <p>I am trying to go through the <a href="http://web.stanford.edu/class/cs224n/assignment1/index.html" rel="nofollow noreferrer">assignment 1</a> for Stanford cs244n class. Problem 1b highly recommend optimization for the Softmax function. I managed to get the Softmax of the N dimensional vector. I also got the Softmax of the MxN dimensional matrix but used a for loop through the columns. I have the following code:</p>
  
  <pre><code>def softmax(x):
      orig_shape = x.shape
  
      # Matrix
      if len(x.shape) &gt; 1:
          softmax = np.zeros(orig_shape)
          for i,col in enumerate(x):
              softmax[i] = np.exp(col - np.max(col))/np.sum(np.exp(col - np.max(col)))
      # Vector
      else:
          softmax = np.exp(x - np.max(x))/np.sum(np.exp(x - np.max(x)))
      return softmax
  </code></pre>
  
  <p>Can I implement a more optimized Matrix implementation?</p>