Overview

PyTorch 1.12, released on June 29, 2022, improves the functional API and extends functorch support.

Main Features

Improved functional API

PyTorch's functional API is enriched with new transformations and better integration with torch.nn.functional.

python 
import torch
import torch.nn.functional as F

# Functional API
x = torch.randn(2, 3, 4)
weights = torch.randn(5, 4)
bias = torch.randn(5)

output = F.linear(x, weights, bias)
print(output.shape)  # torch.Size([2, 3, 5])

# Functional activation and normalization
activated = F.gelu(output)
normalized = F.layer_norm(activated, [5])

Extended functorch

functorch is extended with jacrev and jacfwd for Jacobian computation, and hessian for Hessian matrices.

python 
import torch
from functorch import jacrev, vmap

# Jacobian computation
def f(x):
    return torch.stack([x.sum(), x.prod()])

x = torch.tensor([1.0, 2.0, 3.0])
jacobian = jacrev(f)(x)
print(jacobian.shape)  # (2, 3)

# Batched Jacobian
batch_x = torch.randn(5, 3)
batch_jac = vmap(jacrev(f))(batch_x)
print(batch_jac.shape)  # (5, 2, 3)

Sources