softmax activation

softmax.py

import torch
import torch.nn.functional as F

#----------- Implementing the math -----------#

def softmax(z):
  return z.exp() / z.exp().sum(axis=1, keepdim=True) 
  # keepdim=True tells sum() that we want its output to have the same dimension as z 

zs = torch.tensor([[2., 3., 1.]]) # Three output neurons

sm = softmax(zs)

#----------- Using Pytorch -----------#

torch_sm = F.softmax(zs, dim=1)

#----------- Comparing outputs -----------#

print(f"Pytorch Softmax: \n{torch_sm} \nOur Softmax: \n{sm}")
'''
Out:
Pytorch Softmax: 
tensor([[0.2447, 0.6652, 0.0900]]) 
Our Softmax: 
tensor([[0.2447, 0.6652, 0.0900]])
'''