# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import torch
import torch.nn as nn
from torch.autograd import Variable

# # 1. Does backpropping through a flatten work?

# +
x = Variable(torch.Tensor([1, 2, 3, 4]), requires_grad=True)

z = x.reshape((2, 2))

loss = torch.norm(z) * torch.norm(z)
loss.backward()
# -

x.grad

z

z.grad

x.shape


# # 2. Computing Hessian-vector products

# +
# setup for now:
# f maps [n], [m] to [1], so Hessian is n x n or n x m or m x m depending on what's being done
# Have a vector v, want to compute Hv

def f(x, y):
    return torch.sum(torch.pow(x, 3)) * torch.sum(torch.pow(y, 3))

def g(x, y, theta):
    # return theta**2 * torch.sum(torch.pow(x,3)) * torch.sum(torch.pow(y, 3))
    return x.sum() * y.sum()

n = 4
m = 3
x = torch.Tensor(range(n))
y = torch.Tensor(range(m))
print(f(x, y))
g(x, y, 1)


# -

def hvp(g, x, y, v):
    xvar = Variable(x, requires_grad=True)
    yvar = Variable(y, requires_grad=True)
    vvar = Variable(v, requires_grad=True)
    
    score = g(xvar, yvar)
    
    grad, = torch.autograd.grad(score, yvar, create_graph=True)
    #print(grad)
    total = torch.sum(grad * vvar)
    #print(total)
    
    if xvar.grad:
        xvar.grad.data.zero_()
    if yvar.grad:
        yvar.grad.data.zero_()
        
    total.backward()
    
    if xvar.grad is not None:
        return xvar.grad
    else:
        return torch.zeros(x.shape)


# +
# this will allow us to backpropagate hessian-vector products, the create_graph=True is important

# g is function to take hessian w.r.t.
# x and y are the input variables for g
# x1 and x2 are strings, either 'x' or 'y'
# --- these indicate which variables to take derivatives w.r.t.
# v is vector to get Hessian's action on
def hvp2(g, x, y, x1, x2, v):
    xvar = Variable(x, requires_grad=True)
    yvar = Variable(y, requires_grad=True)
    vvar = Variable(v, requires_grad=True)

    # choose which variable x1var corresponds to
    x1var = xvar if x1=='x' else yvar
    x2var = xvar if x2=='x' else yvar
    
    score = g(xvar, yvar)
    
    grad, = torch.autograd.grad(score, x1var, create_graph=True)
    # print(grad)
    total = torch.sum(grad * vvar)
    # print(total)
    
    if xvar.grad:
        xvar.grad.data.zero_()
    if yvar.grad:
        yvar.grad.data.zero_()
        
    grad2, = torch.autograd.grad(total, x2var, create_graph=True, allow_unused=True)
    # print(grad2)
    return grad2


# -

v = torch.Tensor([2, 1, 0])
print("result", hvp(f, x, y, v))
print("result", hvp2(f, x, y, v))

# +
# so this breaks when the Hessian doesn't survive somehow

# +
# this will not work

# (First time run): comment out following line:
theta.grad.data.zero_()
theta = Variable(torch.Tensor([1]), requires_grad=True)
x = torch.Tensor(range(n))
y = torch.Tensor(range(m))
v = torch.Tensor([2, 1, 0])

score = hvp(lambda x, y: g(x, y, theta), x, y, v)
print(score)
torch.sum(score).backward()
print(theta.grad)

# +
# but this will

# theta.grad.data.zero_()
theta = Variable(torch.Tensor([1]), requires_grad=True)
x = torch.Tensor(range(n))
y = torch.Tensor(range(m))
v = torch.Tensor([2, 1, 0])

score = hvp2(lambda x, y: g(x, y, theta), x, y, v)
print(score)
torch.sum(score).backward()
print(theta.grad)

# +
# \nabla_{xx} g

# theta.grad.data.zero_()
theta = Variable(torch.Tensor([1]), requires_grad=True)
x = torch.Tensor(range(n))
y = torch.Tensor(range(m))
v = torch.Tensor([2, 1, 0, -1])

score = hvp2(lambda x, y: g(x, y, theta), x, y, 'x', 'y', v)
print(score)
torch.sum(score).backward()
print(theta.grad)
# -



StarcoderdataJupyterScriptsDedupFiltered