Skip to content

Latest commit

 

History

History
55 lines (42 loc) · 960 Bytes

File metadata and controls

55 lines (42 loc) · 960 Bytes

Contribution Guidelines

How to write fast PyTorch code

Transpose

torch.transpose(x, -2, -1)
x.transpose(-2, -1)
x.mT   # <-- This is fastest (see #558)

Square

x**2
x * x
torch.square(x)
x.square()   # <-- This is fastest (see #555)

x.op() vs torch.op(x)

torch.any(x < 0)
(x < 0).any()   # <-- This is faster (see #556)
torch.square(x)
x.square()   # <-- This is faster (see #556)
torch.sum(x)
x.sum()   # <-- This is faster (see #556)

Dividing 1 by a tensor

1 / x
torch.reciprocal(x)
x.reciprocal()   # <-- This is fastest (see #563)

Creating new tensors

torch.tensor(0.0, device=a.device, dtype=a.dtype)
torch.zeros((), device=a.device, dtype=a.dtype)
torch.zeros_like(a)   # <-- This is fastest for same shape (see #561)
a.new_zeros(())   # <-- This is fastest for compatible constants (see #561)
a.new_zeros(a.shape)