Skip to content

Instantly share code, notes, and snippets.

View ybj14's full-sized avatar
🏠
Working from home

Binjie Yuan ybj14

🏠
Working from home
  • Tsinghua University
  • Beijing
View GitHub Profile
@ybj14
ybj14 / jvp.py
Last active November 20, 2022 06:51
JVP
# Refer to `https://j-towns.github.io/2017/06/12/A-new-trick.html` for math details.
import torch
from torch import nn
from torch import autograd
def get_jvp(net, x, v):
'''
Generate jacobian vector product. Requires x.requires_grad()
and v.requires_grad().
Args: