Skip to content

Instantly share code, notes, and snippets.

View apaszke's full-sized avatar

Adam Paszke apaszke

View GitHub Profile
import time
import numpy as np
segments = np.ones((100000, 12, 3))
M = np.ones((4, 4))
def _proj_transform_vectors(vecs, M):
is_masked = isinstance(vecs,
if is_masked:
mask = vecs.mask
def patch_matplotlib():
import numpy as np
import time
import matplotlib
from matplotlib import cbook
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import art3d
from mpl_toolkits.mplot3d import proj3d
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.collections import PolyCollection

PyTorch IR

This document presents the IR as of October 17th 2018. Future changes like mutability might render parts of this document outdated.

PyTorch uses an SSA-based IR, which is built of multiple entities:

  • Graph is generally the outermost container for the program representation. At the moment all programs are mostly pure (modulo special operators like prim::Print or prim::PythonOp), but this will change in the future.
  • Then, there are Blocks, which you can treat as functions. They have a list of inputs and outputs, and a (topologically ordered) list of Nodes. Every Graph has a top-level Block which is (using C lingo) like a main function.
  • Nodes, in turn, represent calls to functions (with possibly multiple arguments and returns).
  • Finally, every single intermediate value that appears in the program is represented using a Value - the root Blocks have a list of input values, and every Node takes them as inputs, and returns some more as outputs. Every Value has a Type a
import sys
from collections import OrderedDict
PY2 = sys.version_info[0] == 2
_internal_attrs = {'_backend', '_parameters', '_buffers', '_backward_hooks', '_forward_hooks', '_forward_pre_hooks', '_modules'}
class Scope(object):
def __init__(self):
self._modules = OrderedDict()
import torch
def jacobian(y, x, create_graph=False):
jac = []
flat_y = y.reshape(-1)
grad_y = torch.zeros_like(flat_y)
for i in range(len(flat_y)):
grad_y[i] = 1.
grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
#pragma once
template<typename T>
struct Maybe {
: have_value_(false) {}
Maybe(T value)
: have_value_(true)
, value_(std::move(value)) {}
Maybe(const Maybe&) = delete;
apaszke /
Last active January 16, 2023 07:20
def Rop(y, x, v):
"""Computes an Rop.
y (Variable): output of differentiated function
x (Variable): differentiated input
v (Variable): vector to be multiplied with Jacobian from the right
w = torch.ones_like(y, requires_grad=True)
return torch.autograd.grad(torch.autograd.grad(y, x, w), w, v)
import torch
from torch.autograd import Variable
leaves = [Variable(torch.zeros(5, 5), requires_grad=True) for _ in range(10)]
intermediates = [l + i for i, l in enumerate(leaves)]
loss = sum(v * i for i, v in enumerate(intermediates)).sum()
# define a helper for dividing intermediates into groups
def group(l, group_size):
"""Groups l into chunks of size group_size.
// Writing a string comparison that would consistently work as a compile time
// function, even in older compilers (think GCC 4.8) is not that easy...
template<std::size_t N>
constexpr bool fixed_equal(char const * a, char const * b) {
return *a == *b && fixed_equal<N-1>(a + 1, b + 1);

Managing power saving settings in Linux

Apparently Linux power management can be fairly aggressive and CPU downclocking can have serious effect on PyTorch jobs (even those running on CUDA).

To check the current mode of CPU 0 run the snippet below. On my machine it shows that all cores are in powersave mode.

for CPUFREQ in /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor; do
  [ -f $CPUFREQ ] || continue;
  echo $CPUFREQ $(cat $CPUFREQ)