differentiation of ReLU

relu-diff.py

import torch
import torch.nn.functional as F

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

def relu_prime(z):
  return torch.where(z>0, torch.tensor(1.), torch.tensor(0.))

z = torch.tensor([[-0.2], [0.6]], requires_grad=True)

relu_p = relu_prime(z)

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

torch_relu = F.relu(z)

torch_relu.backward(torch.tensor([[1.], [1.]])) 

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

print(f"Pytorch ReLU': \n{z.grad} \nOur ReLU': \n{relu_p}")
'''
Out:
Pytorch ReLU': 
tensor([[0.],[1.]]) 
Our ReLU': 
tensor([[0.],[1.]])
'''