<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Gradient reversal pytorch
Inspired from the following tweets
Step1:  Tensorflow implementation
Step2:  Modify the gradients
Keep forward pass the same. 
The trick is to add $g(x)$, such that $g'(x)$ is the gradient modifier, during the forward pass and substract it as well. But stop gradients from flowing through the substraction part. 
$f(x) + g(x) - g(x)$ will lead to gradients $f'(x) + g'(x) -g'(x)$. Since gradients don't flow through $-g'(x)$, hence we get new gradients as $f'(x) + g'(x)$
Step3:  Gradient reversal
Here the modifying function $g(x)$ is simply the $-2*f(x)$, this will make the gradients $-f'(x)$.
Step4:  Pytoch case
Step5:  Modify gradients
Step6:  Gradient reversal
Step7:  Pytorch backward hooks
<ASSISTANT_TASK:>
Python Code:
import torch
import tensorflow as tf
from torch.autograd import Variable
import numpy as np
def f(X):
    return X*X
def g(X):
    return X**3
X = np.random.randn(10)
X
Explanation: Gradient reversal pytorch
Inspired from the following tweets:
https://twitter.com/mat_kelcey/status/932149793765261313
https://twitter.com/ericjang11/status/932073259721359363
Basic idea:
```python
Add something to gradient
f(x) + g(x) - tf.stop_gradients(g(x))
Reverse gradient
tf.stop_gradient(f(x)*2) - f(x)
```
End of explanation
sess = tf.InteractiveSession()
tf_X = tf.Variable(X)
init_op = tf.global_variables_initializer()
sess.run(init_op)
sess.run(tf_X)
forward_op = f(tf_X)
sess.run(forward_op)
gradient_op = tf.gradients(forward_op, tf_X)
sess.run(gradient_op)
X*2 # This should match the gradient above
Explanation: Tensorflow implementation
End of explanation
gradient_modifier_op = g(tf_X)
sess.run(gradient_modifier_op)
modified_forward_op = (f(tf_X) + g(tf_X) - tf.stop_gradient(g(tf_X)))
modified_backward_op = tf.gradients(modified_forward_op, tf_X)
sess.run(modified_forward_op)
sess.run(modified_backward_op)
2*X + 3*(X**2) # This should match the gradients above
Explanation: Modify the gradients
Keep forward pass the same. 
The trick is to add $g(x)$, such that $g'(x)$ is the gradient modifier, during the forward pass and substract it as well. But stop gradients from flowing through the substraction part. 
$f(x) + g(x) - g(x)$ will lead to gradients $f'(x) + g'(x) -g'(x)$. Since gradients don't flow through $-g'(x)$, hence we get new gradients as $f'(x) + g'(x)$
End of explanation
gradient_reversal_op = (tf.stop_gradient(2*f(tf_X)) - f(tf_X))
gradient_reversal_grad_op = tf.gradients(gradient_reversal_op, tf_X)
sess.run(gradient_reversal_op)
sess.run(gradient_reversal_grad_op)
sess.run((gradient_op[0] + gradient_reversal_grad_op[0])) # This should be zero. Signifying grad is reversed. 
Explanation: Gradient reversal
Here the modifying function $g(x)$ is simply the $-2*f(x)$, this will make the gradients $-f'(x)$.
End of explanation
def zero_grad(X):
    if X.grad is not None:
        X.grad.data.zero_()
torch_X = Variable(torch.FloatTensor(X), requires_grad=True)
torch_X.data.numpy()
f(torch_X).data.numpy()
g(torch_X).data.numpy()
zero_grad(torch_X)
f_X = f(torch_X)
f_X.backward(torch.ones(f_X.size()))
torch_X.grad.data.numpy()
2*X
Explanation: Pytoch case
End of explanation
modified_gradients_forward = lambda x: f(x) + g(x) - g(x).detach()
zero_grad(torch_X)
modified_grad = modified_gradients_forward(torch_X)
modified_grad.backward(torch.ones(modified_grad.size()))
torch_X.grad.data.numpy()
2*X + 3*(X*X) # It should be same as above
Explanation: Modify gradients
End of explanation
gradient_reversal = lambda x: (2*f(x)).detach() - f(x)
zero_grad(torch_X)
grad_reverse = gradient_reversal(torch_X)
grad_reverse.backward(torch.ones(grad_reverse.size()))
torch_X.grad.data.numpy()
-2*X # It should be same as above
Explanation: Gradient reversal
End of explanation
# Gradient reversal
zero_grad(torch_X)
f_X = f(torch_X)
f_X.register_hook(lambda grad: -grad)
f_X.backward(torch.ones(f_X.size()))
torch_X.grad.data.numpy()
-2*X
# Modified grad example
zero_grad(torch_X)
h = torch_X.register_hook(lambda grad: grad + 3*(torch_X*torch_X))
f_X = f(torch_X)
f_X.backward(torch.ones(f_X.size()))
h.remove()
torch_X.grad.data.numpy()
2*X + 3*(X*X) # It should be same as above
Explanation: Pytorch backward hooks
End of explanation
<END_TASK>