Renormalization Group

Posted on March 1, 2026 by wdengquantum.me

Problem: Consider a Landau-Ginzburg theory involving a single real scalar field $\phi(\mathbf x)$ for $\mathbf x\in\mathbf R^n$ governed by the canonically normalized free energy density:

\[\mathcal F(\phi,\partial\phi/\partial\mathbf x,…)=\frac{1}{2}\biggr|\frac{\partial\phi}{\partial\mathbf x}\biggr|^2+\frac{\phi^2}{2\xi^2}+…\]

Explain what the $+…$ means, explain which terms have temperature $T$-dependence, and explain for which such terms does that $T$-dependence actually matter?

Solution: The $+…$ includes any terms (each with their own coupling constants) compatible with the golden trinity of constraints: locality, analyticity, and symmetry (e.g. a quartic $g\phi^4$ coupling). The part of the free energy density $\mathcal F$ before the $+…$ should compared to the Lagrangian density $\mathcal L$ of Klein-Gordon field theory:

\[\mathcal L=\frac{1}{2c^2}\left(\frac{\partial\phi}{\partial t}\right)^2-\frac{1}{2}\biggr|\frac{\partial\phi}{\partial\mathbf x}\biggr|^2-\frac{\phi^2}{2\bar{\lambda}^2}\]

with $\bar{\lambda}=\hbar/mc$ the reduced Compton wavelength playing a role analogous to the correlation length $\xi\sim 1/\sqrt{|T-T_c|}$; indeed, this $T$-dependence in $\xi=\xi(T)$ is the only $T$-dependence that matters, even though generically all the other coupling constants will also have $T$-dependence.

Problem:

Posted in Blog | Leave a comment

Convolutional Neural Networks

Posted on March 1, 2026 by wdengquantum.me

CNNs_Part_1

$\textbf{Problem}$: Write functions that take an arbitrary grayscale image and convolve them with a Sobel edge detection kernel. Apply both functions to a grayscale image of your choice.

$\textbf{Solution}$:

In [1]:

import numpy as np
import matplotlib.pyplot as plt

sobel_horizontal_kernel = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
sobel_vertical_kernel = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])

def sobel_horizontal(img):
    convolved_image = np.zeros(shape=(img.shape[0]-2,img.shape[1]-2))
    for i in range(1,img.shape[0] - 1):
        for j in range(1,img.shape[1] - 1):
            convolved_image[i-1, j-1] = np.sum(img[i-1:i+2, j-1:j+2] * sobel_horizontal_kernel)
    return convolved_image

def sobel_vertical(img):
    convolved_image = np.zeros(shape=(img.shape[0]-2,img.shape[1]-2))
    for i in range(1,img.shape[0] - 1):
        for j in range(1,img.shape[1] - 1):
            convolved_image[i-1, j-1] = np.sum(img[i-1:i+2, j-1:j+2] * sobel_vertical_kernel)
    return convolved_image

np.random.seed(42)
random_grayscale_img = np.random.randint(0, 256, size=(10, 10))

print("Original Image")
plt.imshow(random_grayscale_img, cmap="gray")
plt.show()

print("Horizontal Sobel")
plt.imshow(sobel_horizontal(random_grayscale_img), cmap="gray")
plt.show()

print("Vertical Sobel")
plt.imshow(sobel_vertical(random_grayscale_img), cmap="gray")
plt.show()

Original Image

No description has been provided for this image

Horizontal Sobel

Vertical Sobel

$\textbf{Problem}$: Above, the $3\times 3$ Sobel edge detection kernels were convolved with an $n\times m$ grayscale image, leading to an $(n-2)\times (m-2)$ edge map. However, although for $n, m \gg 1$ this $-2$ shrinking effect is not particularly noticeable, over many layers of convolutions it can compound undesirably. Thus, one solution is to use $\textit{padding}$ to prevent this shrinking effect. Show how this works.

$\textbf{Solution}$: (by the way, the above unpadded version is sometimes called a $\textit{valid}$ convolution, whereas the padded version is sometimes called a $\textit{same}$ convolution).

In [2]:

padded_random_grayscale_img = np.pad(random_grayscale_img, ((1, 1), (1, 1)), 'constant')
print("Padded random grayscale image shape: ", padded_random_grayscale_img.shape)
plt.imshow(padded_random_grayscale_img, cmap='gray')
plt.show()

print("Horizontal Sobel")
plt.imshow(sobel_horizontal(padded_random_grayscale_img), cmap="gray")
plt.show()

print("Vertical Sobel")
plt.imshow(sobel_vertical(padded_random_grayscale_img), cmap="gray")
plt.show()

Padded random grayscale image shape:  (12, 12)

Horizontal Sobel

Vertical Sobel

$\textbf{Problem}$: Generate a random $15\times 15$ grayscale image, a random $5\times 5$ kernel, and perform a $\textit{stride}$ convolution of the image with the kernel using a stride size of $2$.

$\textbf{Solution}$: The dimension of the output image in this (unpadded) case is $(15-5)/2+1=6$.

In [3]:

np.random.seed(42)
rnd_gray_img = np.random.randint(0, 256, size=(15, 15))
rnd_kernel = np.random.randint(0, 10, size=(5, 5))
plt.imshow(rnd_gray_img, cmap='gray')
plt.show()
plt.imshow(rnd_kernel, cmap='gray')
plt.show()

In [4]:

def stride_convolution(img, kernel, stride_size, padding):
    output_img = np.zeros(shape=(int((img.shape[0] + 2 * padding - kernel.shape[0])/stride_size + 1), int((img.shape[1] + 2 * padding - kernel.shape[1])/stride_size + 1)))
    for i in range(output_img.shape[0]):
        for j in range(output_img.shape[1]):
            output_img[i, j] = np.sum(img[i * stride_size:i * stride_size + kernel.shape[0], j * stride_size:j * stride_size + kernel.shape[1]] * kernel)
    return output_img

plt.imshow(stride_convolution(rnd_gray_img, rnd_kernel, 2, 0), cmap='gray')

Out[4]:

<matplotlib.image.AxesImage at 0x164ca230b90>

$\textbf{Problem}$: Suppose one has a $28\times 28\times 192$ volume at some point inside a convolutional neural network. Explain how a same, unpadded, unit-stride convolution can be performed on this volume to obtain an output volume of dimensions $28\times 28\times 32$.

$\textbf{Solution}$: Use $32$ filters of size $1\times 1\times 192$ with stride $1$ and no padding; this is a common technique to reduce channel number while preserving spatial dimensions.

$\textbf{Problem}$: Implement simplified versions of the LeNet-5, AlexNet, and VGG-16 convolutional neural networks with randomly initialized weights and biases (here simplified means that we will not implement batch normalization, dropout, or any activation functions). Run a forward pass through it without bothering to train the parameters.

$\textbf{Solution}$:

In [18]:

class ConvLayer():
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        
        # Initialize weights and biases
        self.weights = np.random.randn(out_channels, in_channels, kernel_size, kernel_size)
        self.biases = np.random.randn(out_channels)
    
    def forward(self, x):
        # x shape: (batch_size, in_channels, height, width)
        batch_size, in_channels, height, width = x.shape
        
        # Calculate output dimensions
        out_height = (height + 2 * self.padding - self.kernel_size) // self.stride + 1
        out_width = (width + 2 * self.padding - self.kernel_size) // self.stride + 1
        
        # Initialize output
        out = np.zeros((batch_size, self.out_channels, out_height, out_width))
        
        # Apply padding if needed
        if self.padding > 0:
            x = np.pad(x, ((0, 0), (0, 0), (self.padding, self.padding), (self.padding, self.padding)), 'constant')
        
        # Convolution operation
        for b in range(batch_size):
            for c in range(self.out_channels):
                for h in range(out_height):
                    for w in range(out_width):
                        # Extract the receptive field
                        receptive_field = x[b, :, h*self.stride:h*self.stride+self.kernel_size, 
                                            w*self.stride:w*self.stride+self.kernel_size]
                        
                        # Perform convolution
                        out[b, c, h, w] = np.sum(receptive_field * self.weights[c]) + self.biases[c]
        
        return out

class PoolingLayer():
    def __init__(self, kernel_size, stride=2, mode="max"):
        self.kernel_size = kernel_size
        self.stride = stride
        self.mode = mode
    
    def forward(self, x):
        # x shape: (batch_size, channels, height, width)
        batch_size, channels, height, width = x.shape
        
        # Calculate output dimensions
        out_height = height // self.stride
        out_width = width // self.stride
        
        # Initialize output
        out = np.zeros((batch_size, channels, out_height, out_width))
        
        # Max pooling operation
        for b in range(batch_size):
            for c in range(channels):
                for h in range(out_height):
                    for w in range(out_width):
                        # Extract the receptive field
                        receptive_field = x[b, c, h*self.stride:h*self.stride+self.kernel_size, 
                                            w*self.stride:w*self.stride+self.kernel_size]
                        
                        # Perform max pooling
                        if self.mode == "max":
                            out[b, c, h, w] = np.max(receptive_field)
                        elif self.mode == "avg":
                            out[b, c, h, w] = np.mean(receptive_field)
        
        return out

class FCLayer():
    def __init__(self, in_features, out_features):
        self.in_features = in_features
        self.out_features = out_features
        
        # Initialize weights and biases
        self.weights = np.random.randn(out_features, in_features)
        self.biases = np.random.randn(out_features)
    
    def forward(self, x):
        # x shape: (batch_size, in_features)
        batch_size = x.shape[0]
        
        # Initialize output
        out = np.zeros((batch_size, self.out_features))
        
        # Perform matrix multiplication
        for b in range(batch_size):
            out[b] = np.dot(self.weights, x[b]) + self.biases
        
        return out

In [19]:

#Le-Net 5 Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 1, 32, 32))

# Create layers
layer1 = ConvLayer(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0)
layer2 = PoolingLayer(kernel_size=2, stride=2, mode="avg")
layer3 = ConvLayer(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0)
layer4 = PoolingLayer(kernel_size=2, stride=2, mode="avg")
layer5 = FCLayer(in_features=400, out_features=120) # both fully connected layers assume the input image dimensions are (h,w,c) = (32, 32, 1)
layer6 = FCLayer(in_features=120, out_features=84)
layer7 = FCLayer(in_features=84, out_features=1) #y_hat is int b/w 0 and 9 to classify handwritten digit, nowadays prefer softmax

# Forward pass
output = layer1.forward(batch_rand_gray_imgs)
output = layer2.forward(output)
output = layer3.forward(output)
output = layer4.forward(output)
output = output.reshape(num_batches, -1)
output = layer5.forward(output)
output = layer6.forward(output)
output = layer7.forward(output)

print(output)

[[-14004126.31782551]
 [-11998811.82382246]
 [-13235583.18652075]]

In [ ]:

# AlexNet Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 3, 227, 227))

# Create layers
layer1 = ConvLayer(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=0)
layer2 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer3 = ConvLayer(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2) # SAME convolutional layer
layer4 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer5 = ConvLayer(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1)
layer6 = ConvLayer(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1)
layer7 = ConvLayer(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1)
layer8 = PoolingLayer(kernel_size=3, stride=2, mode="max")
layer9 = FCLayer(in_features=9216, out_features=4096) #9216 = 256 * 6 * 6
layer10 = FCLayer(in_features=4096, out_features=4096)
layer11 = FCLayer(in_features=4096, out_features=1000) 

# Forward pass
output = layer1.forward(batch_rand_gray_imgs)
output = layer2.forward(output)
output = layer3.forward(output)
output = layer4.forward(output)
output = layer5.forward(output)
output = layer6.forward(output)
output = layer7.forward(output)
output = layer8.forward(output)
output = output.reshape(num_batches, -1)
output = layer9.forward(output)
output = layer10.forward(output)
output = layer11.forward(output)
# normally would apply softmax on this 1000-feature vector output

for i in range(num_batches):
    plt.semilogy(output[i, :])
    plt.show()

In [32]:

# VGG-16 Implementation
num_batches = 3
np.random.seed(42)
batch_rand_gray_imgs = np.random.randint(0, 256, size=(num_batches, 3, 224, 224))

# Create layers
# Block 1
layer1_1 = ConvLayer(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
layer1_2 = ConvLayer(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)
pool1 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 2
layer2_1 = ConvLayer(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
layer2_2 = ConvLayer(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)
pool2 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 3
layer3_1 = ConvLayer(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)
layer3_2 = ConvLayer(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
layer3_3 = ConvLayer(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)
pool3 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 4
layer4_1 = ConvLayer(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1)
layer4_2 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer4_3 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
pool4 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Block 5
layer5_1 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer5_2 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
layer5_3 = ConvLayer(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1)
pool5 = PoolingLayer(kernel_size=2, stride=2, mode="max")

# Fully Connected Layers
fc1 = FCLayer(in_features=25088, out_features=4096) # 25088 = 512 * 7 * 7
fc2 = FCLayer(in_features=4096, out_features=4096)
fc3 = FCLayer(in_features=4096, out_features=1000)

# Forward pass
output = layer1_1.forward(batch_rand_gray_imgs)
output = layer1_2.forward(output)
output = pool1.forward(output)

output = layer2_1.forward(output)
output = layer2_2.forward(output)
output = pool2.forward(output)

output = layer3_1.forward(output)
output = layer3_2.forward(output)
output = layer3_3.forward(output)
output = pool3.forward(output)

output = layer4_1.forward(output)
output = layer4_2.forward(output)
output = layer4_3.forward(output)
output = pool4.forward(output)

output = layer5_1.forward(output)
output = layer5_2.forward(output)
output = layer5_3.forward(output)
output = pool5.forward(output)

output = output.reshape(num_batches, -1)
output = fc1.forward(output)
output = fc2.forward(output)
output = fc3.forward(output)

for i in range(num_batches):
    plt.semilogy(output[i, :])
    plt.show()

$\textbf{Problem}$: Why is a CNN a better choice of architecture than an MLP for computer vision tasks?

$\textbf{Solution}$: As a simple example, given a convolutional layer of a CNN in which a square $n\times n\times c$ image is convolved (without striding or padding) with $N$ kernels each of size $k\times k$, producing an output volume of dimensions $(n-k+1)\times (n-k+1)\times N$, the number of parameters is expected to be $N(k^2+1)$ (each of the $N$ kernels has $k^2$ weights and a single bias term). By contrast, if one were to instead implement a fully connected layer of dense connections between the flattened input image of size $cn^2$ and the flattened output volume of size $N(n-k+1)^2$, the total number of parameters would be $cn^2N(n-k+1)^2$ (much larger!). Essentially, this is due to the $\textit{parameter sharing}$ property of CNNs, the sparsity of their connections. Finally, the convolutional structure is better at capturing $\textit{translational invariance}$ of the image (unlike the MLP where pixels are just flattened arbitrarily, losing their spatial information).

$\textbf{Problem}$: Explain what a $\textit{residual block}$ is and how several residual blocks may be composed together to define a $\textit{residual neural network}$ (ResNet).

$\textbf{Solution}$: A residual block may be viewed as a simple additive perturbation to a standard $2$-layer feedforward neural network. Recall that, given an input feature vector $\mathbf x$, a feedforward layer $(W,\mathbf b)$ may be viewed as the composition of a $\textit{linear layer}$ $\mathbf x\mapsto W\mathbf x + \mathbf b$ with a nonlinear activation function $\boldsymbol{\sigma}$. Thus, the overall action of a $2$-layer feedforward looks something like: $$\mathbf x\mapsto\boldsymbol{\sigma}(W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2)$$

By contrast, in a residual block, the overall action looks like: $$\mathbf x\mapsto\boldsymbol{\sigma}(W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2 + \mathbf x)$$

(assuming compatible dimensions for vector addition with $\mathbf x$). Thus, although the $1^{\text{st}}$ internal layer of the residual block is no different from a standard feedforward, the difference lies in the insertion of a so-called $\textit{skip connection}$ $+\mathbf x$ in between the linear layer and the activation function of the $2^{\text{nd}}$ layer. A ResNet is then literally just a bunch of residual blocks composed together!

$\textbf{Problem}$: Give some hand-wavy intuition/justification for what makes ResNets interesting/useful.

$\textbf{Solution}$: Consider the following $\textit{thought experiment}$. A standard feedforward neural network classifier with $20$ layers, post-training, has a misclassification rate of $15\%$ $\textit{on the training set}$ itself. Clearly, this seems to be a bias/underfitting problem rather than a variance/overfitting problem because it’s on the $\textit{training set}$. So one reasons that, to make the network more expressive and improve its performance on the training set, one might decide to append another $36$ feedforward layers on top of the $20$ layers already there, thus creating a $56$-layer feedforward. Theoretically, it’s manifest that the optimal training set error can only get lower because one could simply have the first $20$ layers doing exactly what they were doing before, and then have the remaining $36$ layers implement an identity function $\mathbf x\mapsto\mathbf x$. In practice however, it turns out that, paradoxically, training set error actually increases!

The problem is that it turns out to be surprisingly delicate to learn an identity mapping $\mathbf x\mapsto\mathbf x$ due to exploding/vanishing gradient problems during cost function minimization. So rather than struggling so hard to learn the identity, why not redesign the architecture such that the identity $\mathbf x\mapsto\mathbf x$ is the $\textit{default}$ mapping, and to then merely learn whatever perturbation (a.k.a. $\textit{residual}$ hence the network’s name) $\boldsymbol{\varepsilon}(\mathbf x)$ to this identity mapping is needed to learn the actual underlying map of interest $\mathbf x\mapsto \mathbf x+\boldsymbol{\varepsilon}(\mathbf x)$. Roughly speaking, in the earlier notation, one has parameterized the residual by $\boldsymbol{\varepsilon}(\mathbf x)=W_2\boldsymbol{\sigma}(W_1\mathbf x + \mathbf b_1) + \mathbf b_2$, assuming the outer activation function $\boldsymbol{\sigma}$ may be neglected (or if it happens to be a ReLU, this argument is even more appealing!). Notice than that (loosely) $W_1=W_2=\mathbf b_1=\mathbf b_2=\mathbf 0$ gives $\boldsymbol{\varepsilon}(\mathbf x)=\mathbf 0$, hence learning an exact identity map $\mathbf x\mapsto\mathbf x$. But it’s very easy to coerce weights and biases towards $\mathbf 0$ using standard regularization techniques.

In [33]:

# ResNets were first used in computer vision, so the W_1, W_2 matrices above really correspond to convolutional layers in a CNN structure
class ResBlock():
    def __init__(self, in_channels, out_channels, stride=1):
        self.conv1 = ConvLayer(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.conv2 = ConvLayer(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
        
        # Identity x path (Projection Shortcut if dimensions change)
        self.projection = None
        if stride != 1 or in_channels != out_channels:
            # 1x1 convolution to match dimensions
            self.projection = ConvLayer(in_channels, out_channels, kernel_size=1, stride=stride, padding=0)
            
    def relu(self, x):
        return np.maximum(0, x)

    def forward(self, x):
        identity = x
        
        # 1. First transformation + activation
        out = self.conv1.forward(x)
        out = self.relu(out)
        
        # 2. Second transformation
        out = self.conv2.forward(out)
        
        # 3. Match dimensions of identity if needed
        if self.projection is not None:
            identity = self.projection.forward(x)
            
        # 4. Skip Connection Addition
        out += identity
        
        # 5. Final activation
        out = self.relu(out)
        
        return out

$\textbf{Problem}$: Just like a ResNet is made from composing a sequence of residual blocks together, an inception neural network (InceptionNet) is made from composing a sequence of inception blocks together (mostly true; in practice there may also some other side branches, pooling layers, etc). Explain the architecture of an inception block.

$\textbf{Solution}$: Whereas typical CNN architectures require one to choose a filter size, an inception block asks “why not try several filter sizes?”. To this effect, a naive inception block implementation might look something like:

In [34]:

class NaiveInceptionBlock():
    def __init__(self, in_channels, out_channels):
        self.conv1 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=3, stride=1, padding=1) #same convolution
        self.conv3 = ConvLayer(in_channels, out_channels, kernel_size=5, stride=1, padding=2) #same convolution
        self.pool = PoolingLayer(kernel_size=2, stride=2, mode="max")

    def forward(self, x):
        x1 = self.conv1.forward(x)
        x2 = self.conv2.forward(x)
        x3 = self.conv3.forward(x)
        x4 = self.pool.forward(x)
        return np.concatenate([x1, x2, x3, x4], axis=1)

$\textbf{Problem}$: Needless to say, one of the drawbacks to asking “why not try several filter sizes?” is that, by having to literally try several filter sizes, the naive inception block implementation above can get computationally expensive in terms of the number of multiplications and additions required per forward pass. Explain how the introduction of $1\times 1$ $\textit{bottleneck layers}$ can help reduce computational burden without sacrificing too much accuracy.

$\textbf{Solution}$: The idea is to insert $1\times 1$ convolutional layers before the more computationally expensive $3\times 3$ and $5\times 5$ convolutions. These $1\times 1$ convolutions act as “bottlenecks” by reducing the number of channels (and thus the computational cost) before the larger convolutions are applied. For example, instead of applying a $5\times 5$ convolution directly to 256 input channels, we can first apply a $1\times 1$ convolution to reduce the channels to, say, 64, and then apply the $5\times 5$ convolution to the reduced-channel tensor. This significantly reduces the number of parameters and computations while maintaining most of the representational power of the larger filters.

In [35]:

class InceptionBlockWithBottleNeck(NaiveInceptionBlock):
    def __init__(self, in_channels, num_1x1, num_3x3_reduce, num_3x3, num_5x5_reduce, num_5x5, num_pool_proj):
        super().__init__(in_channels, num_1x1, num_3x3, num_5x5, num_pool_proj)
        self.conv_3x3_reduce = ConvLayer(in_channels, num_3x3_reduce, kernel_size=1)
        self.conv_5x5_reduce = ConvLayer(in_channels, num_5x5_reduce, kernel_size=1)
        self.conv_3x3_reduce_out = ConvLayer(num_3x3_reduce, num_3x3, kernel_size=3, padding=1)
        self.conv_5x5_reduce_out = ConvLayer(num_5x5_reduce, num_5x5, kernel_size=5, padding=2)

    def forward(self, x):
        out_1x1 = self.conv_1x1(x)
        out_3x3_reduce = self.conv_3x3_reduce(x)
        out_3x3 = self.conv_3x3_reduce_out(out_3x3_reduce)
        out_5x5_reduce = self.conv_5x5_reduce(x)
        out_5x5 = self.conv_5x5_reduce_out(out_5x5_reduce)
        out_pool = self.pool(x)
        out_pool_proj = self.pool_proj(out_pool)
        return np.concatenate([out_1x1, out_3x3, out_5x5, out_pool_proj], axis=1)

$\textbf{Problem}$: Explain how the MobileNet architecture replaces the standard convolutional layer with a depthwise separable convolutional layer.

$\textbf{Solution}$: The idea is to sort of break down the standard convolutional layer into two parts: a $\textit{depthwise convolutional layer}$ and a $\textit{pointwise convolutional layer}$. The depthwise convolutional layer is a convolutional layer that only convolves with the depth (channels) dimension, while the pointwise convolutional layer is a convolutional layer that only convolves with the height and width dimensions.

Posted in Blog | Leave a comment

Hamilton’s Optics-Mechanics Analogy

Posted on February 19, 2026 by wdengquantum.me

Problem: Deduce the Hamilton-Jacobi equation of classical mechanics.

Solution: Instead of viewing the action $S=S[\mathbf x(t)]$ as a functional of the particle’s trajectory $\mathbf x(t)$, it can be viewed more simply as a scalar field $S(\mathbf x,t)$ in which the initial point in spacetime $(t_0,\mathbf x_0)$ is fixed and one simply takes the on-shell trajectory from $(t_0,\mathbf x_0)$ to $(t,\mathbf x)$. Then the total differential $dS=\mathbf p\cdot d\mathbf x$ (follows from the usual Noetherian calculation) so in particular:

\[\mathbf p=\frac{\partial S}{\partial\mathbf x}\]

Intuitively, this is saying that the particle moves in a direction (the direction of the momentum $\mathbf p$) orthogonal to the contour surfaces of the action field $S$, i.e. such isosurfaces can be viewed as “wavefronts”. Then the total time derivative is:

\[\dot S=L\]

But $\frac{\partial S}{\partial t}+\frac{\partial S}{\partial\mathbf x}\cdot\dot{\mathbf x}=\frac{\partial S}{\partial t}+\mathbf p\cdot\dot{\mathbf x}$. Thus, isolating for $H=\mathbf p\cdot\dot{\mathbf x}-L$ yields the Hamilton-Jacobi nonlinear $1^{\text{st}}$-order PDE for $S(\mathbf x,t)$:

\[-\frac{\partial S}{\partial t}=H\left(\mathbf x,\frac{\partial S}{\partial\mathbf x},t\right)\]

Problem: When $\partial H/\partial t=0$, the Hamiltonian is conserved with energy $H=E$, so this motivates the additive separation of variables, $S(\mathbf x,t):=S_0(\mathbf x)-Et$ for some constant $E$. What does the Hamilton-Jacobi equation simplify to in this case? For a single non-relativistic particle of mass $m$ moving in a potential $V(\mathbf x)$, what does this look like? What about in $1$ dimension?

Solution: \[H\left(\mathbf x,\frac{\partial S_0}{\partial\mathbf x}\right)=E\]

which for $H(\mathbf x,\mathbf p)=|\mathbf p|^2/2m+V(\mathbf x)$ looks like:

\[\frac{1}{2m}\biggr|\frac{\partial S_0}{\partial\mathbf x}\biggr|^2+V(\mathbf x)=E\]

and in $1$ dimension is integrable to the explicit solution:

\[S_0(x)=\pm\int ^xdx’\sqrt{2m(E-V(x’))}\]

In particular, the usual trajectory $x(t)$ can be obtained by treating $S_o=S_0(x,t;E)$ as a family of solutions parameterized by the energy $E$; this works because $S_0$ can be a viewed as a particular generating function of a canonical transformation $(\mathbf x,\mathbf p,H)\mapsto (\mathbf x’,\mathbf p’,H’)$ in which the “boosted” Hamiltonian vanishes $H’=0$.

\[\frac{\partial S_0}{\partial E}=-t_0\Rightarrow t-t_0=\pm\sqrt{\frac{m}{2}}\int_{x_0}^{x(t)}\frac{dx’}{\sqrt{E-V(x’)}}\]

Problem: Above, the static field $S_0(\mathbf x)$ was introduced to simplify the Hamilton-Jacobi equation when the energy $E$ was conserved. However, if one pulls back to the level of functionals rather than fields, one can define an analogous abbreviated action functional $S_0[\mathbf x]$ which depends only on the path $\mathbf x$ taken rather than the trajectory $\mathbf x(t)$. Define $S_0[\mathbf x]$, and moreover show that when the energy $E$ is conserved, the on-shell path is a stationary point of $S_0$ (this is called Maupertuis’s principle).

Solution: The abbreviated action for a single particle of mass $m$ and (non-relativistic) energy $E=|\mathbf p|^2/2m +V(\mathbf x)$ is:

\[S_0[\mathbf x]:=\int d\mathbf x\cdot\mathbf p=\int ds |\mathbf p|=\int ds\sqrt{2m(E-V(\mathbf x))}\]

(modifying this to $S_0[\mathbf x]:=\int ds |\mathbf p|=\int ds\sqrt{2(E-V(\mathbf x))}$ allows for interpretation as an $N$-particle system in configuration space $\mathbf x\in\mathbf R^{3N}$ with the Riemannian “mass metric” $ds^2=m_1|d\mathbf x_1|^2+…+m_N|d\mathbf x_N|^2$).

To find the stationary paths of $S_0[\mathbf x]$ subject to the constraint $H(\mathbf x,\mathbf p)=E$, one can implement a Lagrange multiplier $\gamma(\tau)$ to perform unconstrained extremization of:

\[S[\mathbf x(\tau)]:=S_0[\mathbf x(\tau)]-\int d\tau\gamma(\tau)(H(\mathbf x,\mathbf p)-E)=\int d\tau (\mathbf p\cdot\dot{\mathbf x}-\gamma(\tau)(H(\mathbf x,\mathbf p)-E))\]

The Euler-Lagrange equations lead to Hamilton’s equations:

\[\frac{d\mathbf x}{d\tau}=\gamma\frac{\partial H}{\partial\mathbf p}\]

\[\frac{d\mathbf p}{d\tau}=-\gamma\frac{\partial H}{\partial\mathbf x}\]

provided the Lagrange multiplier $\gamma=dt/d\tau$ encodes reparameterization invariance; with this choice it’s clear that the integrand in the functional $S$ was nothing more than the Lagrangian $L=\mathbf p\cdot\dot{\mathbf x}-H$ (plus an unimportant constant $E$) so Maupertuis’s principle reduces to the usual Hamilton’s principle.

Problem: What does Fermat’s principle in ray optics assert? Hence, derive the ray equation.

Solution: The time functional $T=T[\mathbf x(s)]$ of a ray trajectory $\mathbf x(s)$ is stationary on-shell. That is:

\[cT[\mathbf x(s)]=\int ds n(\mathbf x(s))\]

This is reparameterization invariant, since one can arbitrarily parameterize $\mathbf x=\mathbf x(t)$ and replace $ds=dt|\dot{\mathbf x}|$. The corresponding Euler-Lagrange equations are:

\[\frac{d}{dt}\left(n(\mathbf x)\frac{\dot{\mathbf x}}{|\dot{\mathbf x}|}\right)=|\dot{\mathbf x}|\frac{\partial n}{\partial\mathbf x}\]

But by choosing the natural parameterization $t:=s$ one has $|d\mathbf x/ds|=1$, hence the ray equation:

\[\frac{d}{ds}\left(n\frac{d\mathbf x}{ds}\right)=\frac{\partial n}{\partial\mathbf x}\]

This can also be written in terms of the curvature vector $\boldsymbol{\kappa}=d^2\mathbf x/ds^2$:

\[\boldsymbol{\kappa}=\left(\frac{\partial\ln n}{\partial\mathbf x}\right)_{\perp d\mathbf x}\]

Problem: Starting from an arbitrary Cartesian component $\psi(\mathbf x,t)=\psi_0(\mathbf x)e^{i(k_0cT(\mathbf x)-\omega t)}$ of either the $\mathbf E$ or $\mathbf B$ fields of a light wave (here $\omega=ck_0$ with $k_0=2\pi/\lambda_0$ is the free space wavenumber), make the eikonal approximation to the dispersionless wave equation obeyed by $\psi$ in order to obtain the (scalar) eikonal equation. By defining light rays as the integral curves of the eikonal field $cT(\mathbf x)$ (a kind of local optical path length), reproduce the vector eikonal equation from Fermat’s principle above.

Solution: The ansatz $\psi(\mathbf x,t)=\psi_0(\mathbf x)e^{i(k_0cT(\mathbf x)-\omega t)}$ is easy to justify; the $e^{-i\omega t}$ is a just a Fourier transform factor that reduces the wave equation $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi=\frac{n^2}{c^2}\frac{\partial^2\psi}{\partial t^2}$ to a Helmholtz equation $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi=-n^2k_0^2\psi$. The remaining piece is just a polar parameterization of an arbitrary $\mathbf C$-valued spatial field $\psi_0(\mathbf x)e^{ik_0cT(\mathbf x)}$. One obtains:

\[\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|^2=n^2+\frac{1}{k_0^2\psi_0}\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\psi_0+\frac{2i}{k_0\psi_0}\frac{\partial\psi_0}{\partial\mathbf x}\cdot\frac{\partial cT}{\partial\mathbf x}+\frac{i}{k_0}\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2cT\]

The eikonal approximation amounts to taking the ray optics limit $k_0\to\infty$ (in practice, the wavelength $2\pi/k_0$ has to be much shorter than all other length scales), and yields the (scalar) eikonal equation:

\[\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|=n\]

A light ray is thus a trajectory $\mathbf x(s)$ with unit tangent vector:

\[\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial cT}{\partial\mathbf x}\]

The rest is an application of the chain rule:

\[\frac{d}{ds}=\frac{\partial}{\partial\mathbf x}\cdot\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial}{\partial\mathbf x}\cdot\frac{\partial cT}{\partial\mathbf x}\]

followed by the identity:

\[\left(\frac{\partial cT}{\partial\mathbf x}\cdot\frac{\partial}{\partial\mathbf x}\right)\frac{\partial cT}{\partial\mathbf x}=\frac{1}{2}\frac{\partial}{\partial\mathbf x}\biggr|\frac{\partial cT}{\partial\mathbf x}\biggr|^2\]

to deduce the (vector) eikonal equation of motion for ray trajectories just as Fermat’s principle predicts.

Problem: Hence, what is Hamilton’s optics-mechanics analogy?

Solution: In a nutshell, the isomorphism proceeds as:

\[(n(\mathbf x), cT)\leftrightarrow (|\mathbf p(\mathbf x)|,S_0)\]

Problem: Use Hamilton’s optics-mechanics analogy to solve the brachistochrone problem (this was how Johann Bernoulli originally solved it).

Solution: By energy conservation, the speed of the particle at distance $y>0$ below its initial dropping height is $v=\sqrt{2gy}$. By Fermat’s principle, minimizing the time functional then amounts to treating the particle as a light ray with $n(\mathbf x)=c/v(y)$. So the question becomes how do light rays bend in a horizontally stratified medium with $n(y)\propto y^{-1/2}$? The answer is given by the ray equations:

\[\frac{d}{ds}\begin{pmatrix} y^{-1/2}dx/ds \\ y^{-1/2}dy/ds\end{pmatrix}=\begin{pmatrix}0 \\ y^{-1/2}/2\end{pmatrix}\]

The horizontal component expresses Snell’s law since $dx/ds=\sin\theta$ (it expresses momentum conservation along the homogeneous $\partial n/\partial x=0$ direction). Using the tangent vector constraint $ds^2=dx^2+dy^2$ gives the ODE of a cycloid:

\[\frac{dy}{dx}=\sqrt{\frac{\text{const}}{y}-1}\]

(the vertical component ODE has an analytical solution $y(s)=-s^2/8R+s$ which is contained in the cycloid, so is redundant information).

Problem: How did Hamilton’s optics-mechanics analogy inspire Schrodinger to propose his famous equations of quantum mechanics?

Solution: Essentially, Schrodinger asked: ray optics is to wave optics as classical mechanics is to what? In other words, one imagines there exists a wave theory of particles/matter and one would like to take the “inverse eikonal limit” of classical mechanics (here, “inverse eikonal limit” is usually called quantization):

Just as light rays propagate parallel to their phase fronts:

\[\frac{d\mathbf x}{ds}=\frac{1}{n}\frac{\partial cT}{\partial\mathbf x}\]

Particles propagate parallel to their “action fronts” exactly according to Hamilton’s analogy:

\[\frac{d\mathbf x}{ds}=\frac{1}{|\mathbf p(\mathbf x)|}\frac{\partial S_0}{\partial\mathbf x}\]

Already this suggests that the action should have some phase interpretation. More precisely, it should be the phase of the particle’s de Broglie wave in units of $\hbar$. It’s also not obvious that particle’s should be described by a scalar wave field rather than e.g. the electromagnetic vector wave fields of light. Schrodinger simply guessed it looked like the equivalent of “scalar diffraction theory” with a single wavefunction $\psi(\mathbf x,t)=\psi_0(\mathbf x,t)e^{iS(\mathbf x,t)/\hbar}$. This gives the Madelung equations of quantum hydrodynamics, one of which is just a continuity equation (giving credence to the Born interpretation of $|\psi|^2$) and the other is a quantum Hamilton-Jacobi equation which in the limit $\hbar\to 0$ (analogous to the eikonal limit $\lambda_0\to 0$) simplifies to the classical Hamilton-Jacobi equation.

Posted in Blog | Leave a comment

Pseudo-Riemannian Geometry

Posted on February 13, 2026 by wdengquantum.me

Problem: Define the signature of a matrix. Hence, state and prove Sylvester’s law of inertia.

Solution: The signature of an $n\times n$ matrix $A$ is a $3$-tuple $(n_+,n_-,n_0)$ where $n_+$ is the number of positive eigenvalues of $A$ (including multiplicity), $n_-$ is the number of negative eigenvalue of $A$ (including multiplicity), and $n_0=\text{dim}\ker A$ is the multiplicity of the zero eigenvalue; thus, $n_++n_-+n_0=n$.

Let $A,B$ be real symmetric matrices. Then Sylvester’s law of inertia asserts that $A$ and $B$ are congruent matrices iff they have the same signature (which sometimes is called “inertia” because of this invariance under congruence, hence the name).

Proof: It’s easy to see that the nullity $n_0$ is preserved by congruence transformations. If one can can show that $n_+$ is also preserved, then it implies $n_-$ is conserved by virtue of $n_++n_-+n_0=n$. To show this, the idea is to prove by contradiction, assuming $n_+$ is not preserved which by dimension counting would imply a non-zero vector living in the intersection of the subspace spanned by the positive-eigenvalue eigenvectors of $A$ and the congruence-transformed subspace spanned by the non-positive-eigenvalue eigenvectors of $B$.

Since any real, symmetric matrix $A$ is isomorphic to a real quadratic form $Q(\mathbf x):=\mathbf x^TA\mathbf x$, the concepts of signature and Sylvester’s law of inertia can also be reformulated in the language of quadratic forms rather than real symmetric matrices.

Problem: Let $X$ be a smooth $n$-manifold. Explain what it means to place the additional structure of a pseudo-Riemannian metric $g$ on $X$. Then explain how Riemannian and Lorentzian geometry are special cases of pseudo-Riemannian geometry.

Solution: A Riemannian metric $g$ on $X$ is a type $(0,2)$ tensor field that defines an inner product $g_x:T_x(X)^2\to\mathbf R$ at the tangent space $T_x(X)$ of each point $x\in X$. The “tensor field” part is sometimes rephrased as saying that $g_x$ is a bilinear form. Moreover, to flesh out the usual axioms of a real inner product space, the Riemannian metric tensor $g_x$ at each $x\in X$ must be symmetric $g_x(v_x,v’_x)=g_x(v’_x,v_x)$ and positive-definite $v_x\neq 0\Rightarrow g_x(v_x,v_x)>0$.

A pseudo-Riemannian metric relaxes the positive-definite requirement of a Riemannian metric, instead merely requiring non-degeneracy (i.e. at each point $x\in X$, the zero vector $0$ is the only vector orthogonal $g_x(0,v_x)=0$ to all $v_x\in T_x(X)$). More concisely, what this is saying is that $n_0=0$, so the signature of a pseudo-Riemannian metric may be thought of as a pair $(n_+,n_-)$.

Thus, Riemannian metrics are the special subset of pseudo-Riemannian metrics for which $(n_+,n_-)=(n,0)$. Meanwhile, Lorentzian metrics are another special subset of pseudo-Riemannian metrics for which $(n_+,n_-)=(n-1,1)$ (this is the relativists’/ convention). Typically, spacetime $X$ has dimension $n=4$, so e.g. the Minkowski metric is said to have signature $(3,1):=(-,+,+,+)$.

Problem: Let $(X,g)$ be a pseudo-Riemannian manifold. Often, one writes the general expression for an infinitesimal line element on $X$ as $ds^2=g_{\mu\nu}dx^{\mu}dx^{\nu}$; explain the shorthand being used here.

Solution: This is nothing more than choosing some chart $x^{\mu}$ on $X$ and expanding the metric tensor field $g$ in the coordinate basis $\{dx^{\mu}\otimes dx^{\nu}\}$ of type $(0,2)$ tensor fields:

\[g=g_{\mu\nu}dx^{\mu}\otimes dx^{\nu}\]

for some real, symmetric scalar fields $g_{\mu\nu}:X\to\mathbf R$. One then simply writes $ds^2:=g$ and omits the tensor product $\otimes$.

Problem: Let $x(t)\in X$ be a curve on a Riemannian manifold $(X,g)$. By choosing a chart $x^{\mu}$ to cover a suitable region of $X$ spanned by the trajectory $x(t)$, explain how to compute the length $s$ of the curve $x(t)$ using the Riemannian metric $g$.

Solution: Heuristically, it is:

\[s=\int ds=\int\sqrt{g_{\mu\nu}dx^{\mu}dx^{\nu}}=\int dt\sqrt{g_{\mu\nu}(x(t))\dot{x}^{\mu}\dot{x}^{\nu}}\]

Problem: Write down the action $S$ for a non-relativistic particle of mass $m$ moving on a Riemannian manifold $(X,g)$. Hence, derive the geodesic equation. What happens if the particle is relativistic?

Solution: A nonrelativistic free particle of mass $m$ has action $S=\int dt L$ described by the usual nonrelativistic kinetic Lagrangian:

\[L=\frac{m}{2}|\dot{\mathbf x}|^2\]

The catch is that $|\dot{\mathbf x}|^2=g_{ij}(\mathbf x)\dot x^{i}\dot x^{j}$ depends on the Riemannian metric $g$ in which the particle moves. Applying the Euler-Lagrange equations, one finds the equation of motion (called the geodesic equation):

\[\ddot x^i+\Gamma^i_{jk}\dot x^j\dot x^k=0\]

where the Christoffel symbols $\Gamma^i_{jk}(\mathbf x):=\frac{1}{2}g^{i\ell}(\mathbf x)\left(\frac{\partial g_{\ell j}}{\partial x^k}+\frac{\partial g_{\ell k}}{\partial x^j}-\frac{\partial g_{jk}}{\partial x^{\ell}}\right)$ are analogous to “fictitious forces” on the particle due to the curvature of $X$.

By contrast, a relativistic free particle has action $S=-mcs=-mc^2\tau$. This can be recast into non-relativistic form $S=\int dt L$ using the relativistic kinetic Lagrangian:

\[L=-mc\sqrt{g_{\mu\nu}\dot x^{\mu}\dot x^{\nu}}\]

which in flat space $g_{\mu\nu}=\text{diag}(1,-1,-1,-1)$ reduces to the familiar $L=-mc^2/\gamma$ (nb. more generally, the quantity under the square root is always $\geq 0$ for the timelike worldlines of particles). Some comments:

The stationary action principle $\delta S=0$ in which the action $S$ tends to be minimized corresponds to the well-known fact that relativistic free particles travelling in straight line geodesics through spacetime maximize the Minkowski distance $s$, or equivalently Alice experiences more proper time (i.e. ages more) than Bob in the twin paradox.
The relativistic action (unlike its non-relativistic counterpart) is manifestly reparameterization invariant.

Now the Euler-Lagrange geodesic equations for the relativistic Lagrangian are almost identical to the non-relativistic case except there’s an additional term:

\[\ddot x^{\mu}+\Gamma^{\mu}_{\nu\rho}(x)\dot x^{\nu}\dot x^{\rho}=\frac{\dot L}{L}\dot x^{\mu}\]

The additional term vanishes $\dot L=0$ iff $L=L(\tau’)$ is parameterized as any affine transformation $\tau’=a\tau+b$ of the particle’s proper time $\tau$ (prove this by showing $\dot{\tau}’=-aL/mc^2$). Only in this case does the relativistic geodesic equation reduce to the non-relativistic form:

\[\frac{d^2 x^{\mu}}{d\tau’^2}+\Gamma^{\mu}_{\nu\rho}(x(\tau’))\frac{dx^{\nu}}{d\tau’}\frac{dx^{\rho}}{d\tau’}=0\]

Problem: For $X=\mathbf R^2$, calculate all non-vanishing Christoffel symbols in the polar coordinate chart $(\rho,\phi)$.

Solution:

Note it is usually more efficient to compute Christoffel symbols this way rather than using the explicit formula in terms of partial derivatives of the metric $g$ (though they’re of course equivalent).

Problem: Explain how the presence of a pseudo-Riemannian metric $g$ on a smooth manifold $X$ allows one to identify covectors in $T_x^*(X)$ with tangent vectors in $T_x(X)$ at each point $x\in X$.

Solution: Just as in $X=\mathbf R^n$, one identifies a tangent vector $\mathbf v$ with its covector $\mathbf v\cdot$ via the inner product $\cdot$, so at a given $x\in X$, one uses the pseudo-inner product $g_x$ on $T_x(X)$ to make the exact same identification $v_x\leftrightarrow g_x(v_x,\space)$ (this sometimes called the musical isomorphism induced by $g$ in light of the map $\flat_x:T_x(X)\to T^*_x(X)$ defined by $v_x^{\flat_x}(v’_x):=g(v_x,v’_x)$ and its inverse $\sharp_x=\flat_x^{-1}$).

In some chart $x^{\mu}$, one can write the $1$-form $v^{\flat}=(v^{\flat})_{\mu}dx^{\mu}$ with components $(v^{\flat})_{\mu}=v^{\flat}(\partial_{\mu})=g(v,\partial_{\mu})=g_{\mu\nu}v^{\nu}$ where $v^{\nu}=v(x^{\nu})$ are the components of its musically isomorphic vector field $v=v^{\nu}\partial_{\nu}$ in the same coordinate basis. Physicists typically abuse notation by merely writing $(v^{\flat})_{\mu}=g_{\mu\nu}v^{\nu}$ as just $v_{\mu}=g_{\mu\nu}v^{\nu}$, identifying $v^{\flat}\equiv v$ under the musical isomorphism. But then this lazy notation makes it look as if the metric $g$ performs a trivial mechanical action of just “lowering the index” on $v$.

Problem: Define a tensor field $g^*$ that performs the trivial mechanical action of “raising the index” on a $1$-form $A$.

Solution: Define $g^*_x:T^*_x(X)^2\to\mathbf R$ to be a type $(2,0)$ tensor given by:

\[g^*(A,A’):=g(A^{\sharp}, A’^{\sharp})\]

Then on the one hand, in some chart $x^{\mu}$, one has:

\[g^*(v^{\flat},v’^{\flat})=g(v,v’)=g_{\mu\nu}v^{\mu}v’^{\nu}\]

On the other hand, in that same chart $x^{\mu}$, one has:

\[g^*(v^{\flat},v’^{\flat})=(g^*)^{\mu\nu}(v^{\flat})_{\mu}(v’^{\flat})_{\nu}=(g^*)^{\mu\nu}g_{\mu\rho}v^{\rho}g_{\nu\sigma}v’^{\sigma}=(g^*)^{\rho\sigma}g_{\mu\rho}g_{\nu\sigma}v^{\mu}v’^{\nu}\]

This implies $g_{\mu\nu}=(g^*)^{\rho\sigma}g_{\mu\rho}g_{\nu\sigma}$. Since $g$ is non-degenerate (the $N_0=0$ axiom of a pseudo-Riemannian metric), it follows that $\det g_{\mu\nu}\neq 0$ is invertible, so the only possibility is $(g^*)^{\mu\nu}g_{\nu\rho}=\delta^{\mu}_{\rho}$, i.e. $g^*\sim g^{-1}$. It is common practice to just drop the $*$ and write $g^{\mu\nu}$ in lieu of $(g^*)^{\mu\nu}$. With this lazy notation, $g^{\mu\nu}$ looks like it just “raises the index” on $A^{\mu}=g^{\mu\nu}A_{\nu}$ (where one last abuse of notation $(A^{\sharp})^{\mu}=A^{\mu}$ has been made!). cf. raising lowering operators $a,a^{\dagger}$ in quantum mechanics.

Problem: What is the canonical volume form on a smooth, orientable pseudo-Riemannian $n$-manifold $X$? Why is it canonical?

Solution: In a chart $x^{\mu}$ in which the pseudo-Riemannian metric has components $g=g_{\mu\nu}dx^{\mu}\otimes dx^{\nu}$, the canonical volume form is $\sqrt{|\det g|}dx^1\wedge…\wedge dx^n$. This is indeed a volume form because it is a top form with $\det g\neq 0$ nowhere vanishing for same reasons as before. Moreover, it is canonical because it does not depend on the choice of chart $x^{\mu}$ (provided it’s the same orientation) as $\sqrt{|\det g|}$ is a scalar density of weight $1$ whereas $dx^1\wedge…\wedge dx^n$ is a scalar density of weight $-1$.

Problem: Let $X$ be a smooth, $n$-dimensional, orientable pseudo-Riemannian manifold, and let $\omega\in\Omega^k(X)$ be a differential $k$-form on $X$. Define the Hodge dual $n-k$-form $\star\omega\in\Omega^{n-k}(X)$ on $X$.

Solution: The Hodge dual $\star\omega\in\Omega^{n-k}(X)$ is defined to be the unique differential $n-k$-form such that …

Posted in Blog | Leave a comment

Reinforcement Learning (Part $1$)

Posted on February 1, 2026 by wdengquantum.me

Problem: How does the paradigm of reinforcement learning (RL) fit within the broader context of machine learning?

Solution: It is instructive to compare/contrast reinforcement learning with supervised learning. In this way, it will be seen that RL can in fact be viewed as a generalization of SL:

In SL, the training feature vectors $\mathbf x_i$ manifest as the state vector $\mathbf x_t$ of an agent (similar to how it’s okay to conflate the position vector $\mathbf x(t)$ of a particle in classical mechanics with the particle itself, there’s no harm in conflating the state vector $\mathbf x_t$ with the agent itself, see map-territory relation). The difference is purely a perspective shift reflected in the choice of subscript variable $i=1,2,3,…$ vs. $t=0,1,2,…$.
In SL, the feature vectors $\mathbf x_i$ are i.i.d. draws from the same underlying data-generating distribution. In RL, the initial state vector $\mathbf x_{t=0}$ is drawn from an initial distribution, and subsequent state vectors $\mathbf x_t$ are dependent on the previous state $\mathbf x_{t-1}$ (discrete-time Markov chain).
In SL, the model output $\hat y(\mathbf x|\boldsymbol{\theta})$ is analogous to an action $\mathbf f$ outputted by an agent’s policy function $\pi(\mathbf f|\mathbf x,t)$.
In SL, each training feature vector $\mathbf x_i$ is labelled by its correct output $y_i$. This concept has no analog in reinforcement learning as there is no notion of a “correct” action for the agent to take.
In SL, the feedback mechanism is provided by a loss function $L(\hat y,y)$. By contrast, in RL the feedback mechanism is provided by a reward signal $r_t$.
In SL, the objective is to find optimal parameters $\boldsymbol{\theta}^*=\text{argmin}_{\boldsymbol{\theta}}\sum_{i=1}^{N_{\text{train}}}L(\hat y(\mathbf x_i|\boldsymbol{\theta}),y_i)$ minimizing the training cost over the $N_{\text{train}}$ training examples. In RL, the objective is to find an optimal agent policy $\pi^*=\text{argmax}_{\pi}\langle\sum_{t=0}^T\gamma^tr_t|\pi\rangle$ maximizing the expected return over an episode of horizon $T\leq\infty$ (which depends on if/when the agent reaches a terminal state $\mathbf x_T$).
In SL, there is often a distinction between training, cross-validation, and testing datasets, with a golden rule often being that the model should obviously not be able to see any of the testing data. In RL, it is societally acceptable to train on your test set 🙂

Problem: Imagine the set of all states $\mathbf x$, i.e. the agent’s state space. On this state space, one can impose a scalar field $v_{\pi}(\mathbf x,t)$; what is the meaning of this field? Give a simple example thereof.

Solution: The scalar field $v_{\pi}(\mathbf x,t)$ can roughly be thought of as describing how “valuable” the state $\mathbf x$ is at time $t$. More precisely, if one were to initialize an agent $\pi$ at state $\mathbf x_t:=\mathbf x$, then $v_{\pi}(\mathbf x,t)$ is the expected return:

\[v_{\pi}(\mathbf x,t):=\biggr\langle\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\biggr|\mathbf x,\pi\biggr\rangle\]

Consider a $\pi$-creature (of $3$Blue$1$Brown fame) symbolizing an agent with policy $\pi$. This $\pi$-creature has to complete a maze. Within the RL framework, this can be modelled as an MDP in which each square is one possible state $\mathbf x$ of the agent, the maze structure defines the allowed actions $\mathbf f$ from each state $\mathbf x$, and one can work with an undiscounted $\gamma=1$ return in which the reward is $r_t=-1$ for each action taken. Thus, the optimal policy $\pi^*$ is a time-independent, deterministic policy $\mathbf f=\pi^*(\mathbf x)$ that gets the $\pi$-creature from its initial state to the terminal state in the fewest number of moves. With that in mind, one can label on top of each square $\mathbf x$ its corresponding optimal value $v_{\pi^*}(\mathbf x)$ (note, this image was generated using Nano Banana Pro, some of the calculated values are just wrong but the idea is clear):

Problem: The value function $v_{\pi}(\mathbf x,t)$ is passive; it simply assess how good/bad it is to be at $\mathbf x_t:=\mathbf x$. In order to remedy this, one can look at the quality function $q_{\pi}(\mathbf x,\mathbf f,t)$; explain how this takes on a more active role compared to the passive nature of the value function $v_{\pi}(\mathbf x,t)$.

Solution: Because $q_{\pi}(\mathbf x,\mathbf f,t)$ assesses the quality of the agent choosing the action $\mathbf f_t:=\mathbf f$ while in state $\mathbf x_t:=\mathbf x$. That is, it is a more refined conditional expected return:

\[q_{\pi}(\mathbf x,\mathbf f,t)=\biggr\langle\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\biggr|\mathbf x,\mathbf f,\pi\biggr\rangle\]

Problem: How come the agent’s policy, value function and quality function are sometimes respectively written as $\pi(\mathbf f|\mathbf x),v_{\pi}(\mathbf x)$ and $q_{\pi}(\mathbf x,\mathbf f)$ without the $t$ argument?

Solution: There could be $3$ reasons:

Similar to the above maze example, sometimes it just doesn’t depend on $t$.
If one is seeking to maximize the agent’s expected return over an infinite horizon $T=\infty$; in such a case, it is both intuitively and mathematically clear that $\pi(\mathbf f|\mathbf x),v_{\pi}(\mathbf x)$ and $q_{\pi}(\mathbf x,\mathbf f)$ should all be invariant with respect to time translations.
If the horizon is finite $T<\infty$, then it is standard to package $\mathbf x$ and $t$ together, thereby working with an “augmented” state vector $\mathbf x’:=(\mathbf x,t)$. Then everything can be made to look identical to the infinite horizon $T=\infty$ case by replacing $\mathbf x\mapsto\mathbf x’$.

(cf. distinction between state variables and path variables in thermodynamics).

Problem: State and prove the law of iterated expectation.

Solution: If $X,Y$ are any random variables, then:

\[\langle X\rangle=\langle\langle X|Y\rangle\rangle\]

The proof is easy as long as one is careful to interpret all the expectations correctly. For instance, $\langle X|Y\rangle$ is not a scalar but a random variable with respect to $Y$:

\[\langle X|Y\rangle=\int dx p(x|Y) x\]

Thus, it is clear that the outer expectation must also be with respect to $Y$ alone:

\[\langle\langle X|Y\rangle\rangle=\int dy p(y)\langle X|y\rangle=\int dxdy p(y)p(x|y)x\]

Rewriting in terms of the joint distribution $p(x,y)=p(y)p(x|y)$ and integrating out $\int dy p(x,y)=p(x)$, one finally obtains:

\[=\int dx p(x) x=\langle X\rangle\]

Of course, this also generalizes easily to identities such as the following equality of random variables (w.r.t. $Z$):

\[\langle X|Z\rangle=\langle\langle X|Y,Z\rangle|Z\rangle\]

Problem: Show that the value function $v_{\pi}(\mathbf x)$ satisfied the following identity (to be used as a lemma later):

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

Solution: Denoting the return random variable by:

\[R_t:=\sum_{t’=t+1}^T\gamma^{t’-t-1}r_{t’}\]

The fundamental observation is that $R_t$ obeys a simple recurrence relation:

\[R_t=r_{t+1}+\gamma R_{t+1}\]

Taking suitable conditional expectations thereof:

\[\langle R_t|\mathbf x_t,\pi\rangle=\langle r_{t+1}|\mathbf x_t,\pi\rangle+\gamma\langle R_{t+1}|\mathbf x_t,\pi\rangle\]

The LHS is the definition of $v_{\pi}(\mathbf x_t)$. Applying the law of iterated expectation to the $2^{\text{nd}}$ term on the RHS:

\[\langle R_{t+1}|\mathbf x_t,\pi\rangle=\langle\langle R_{t+1}|\mathbf x_{t+1},\mathbf x_t,\pi\rangle|\mathbf x_t,\pi\rangle\]

The Markov property ensures that $\langle R_{t+1}|\mathbf x_{t+1},\mathbf x_t,\pi\rangle=\langle R_{t+1}|\mathbf x_{t+1},\pi\rangle=v_{\pi}(\mathbf x_{t+1})$. One thus obtains the desired result:

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

Problem: How are $v_{\pi}(\mathbf x_t)$ and $q_{\pi}(\mathbf x_t,\mathbf f_t)$ related to each other?

Solution: Apply the law of iterated expectations:

\[v_{\pi}(\mathbf x_t):=\langle R_t|\mathbf x_t,\pi\rangle=\langle\langle R_t|\mathbf x_t,\mathbf f_t,\pi\rangle|\mathbf x_t,\pi\rangle=\langle q_{\pi}(\mathbf x_t,\mathbf f_t)|\mathbf x_t,\pi\rangle\]

Since $\mathbf x_t$ is being conditioned upon, the expectation must be over $\mathbf f_t$ so one can explicitly write it as:

\[=\sum_{\mathbf f_t}p(\mathbf f_t|\mathbf x_t,\pi)q_{\pi}(\mathbf x_t,\mathbf f_t)\]

But trivially $p(\mathbf f_t|\mathbf x_t,\pi)=\pi(\mathbf f_t|\mathbf x_t)$. Thus, one has an expression for $v_{\pi}(\mathbf x_t)$ in terms of $q_{\pi}(\mathbf x_t,\mathbf f_t)$. Now one would like to proceed the other way, relating $q_{\pi}(\mathbf x_t,\mathbf f_t)$ to $v_{\pi}(\mathbf x_t,\mathbf f_t)$. This can be achieved by fleshing out the expectation from the earlier lemma:

\[v_{\pi}(\mathbf x_t)=\langle r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1})|\mathbf x_t,\pi\rangle\]

\[=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\pi)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

Further expand $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\pi)=\sum_{\mathbf f_t}p(\mathbf f_t|\mathbf x_t,\pi)p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t,\pi)$ and recognize again $p(\mathbf f_t|\mathbf x_t,\pi)=\pi(\mathbf f_t|\mathbf x_t)$ and $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t,\pi)=p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)$ because the action $\mathbf f_t$ is already fixed. Moving the summation $\sum_{\mathbf f_t}$ to the outside so as to compare with the earlier identity expressing $v_{\pi}(\mathbf x_t)$ in terms of $q_{\pi}(\mathbf x_t,\mathbf f_t)$, one can pattern-match:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

(which in hindsight is manifest…)

Problem: Hence, deduce the Bellman equations for the value and quality functions.

Solution: To obtain the Bellman equation for $v_{\pi}(\mathbf x_t)$, substitute the above expression for $q_{\pi}(\mathbf x_t,\mathbf f_t)$ into the expression for $v_{\pi}(\mathbf x_t)$:

\[v_{\pi}(\mathbf x_t)=\sum_{\mathbf f_t}\pi(\mathbf f_t|\mathbf x_t)\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi}(\mathbf x_{t+1}))\]

Similarly, to the get the Bellman equation for $q_{\pi}(\mathbf x_t,\mathbf f_t)$ substitute vice versa:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)\left(r_{t+1}+\gamma\sum_{\mathbf f_{t+1}}p(\mathbf f_{t+1}|\mathbf x_{t+1},\pi)q_{\pi}(\mathbf x_{t+1},\mathbf f_{t+1})\right)\]

Intuitively, the Bellman equations relate the value of every state with the values of states that can be transitioned into, thereby providing a system of simultaneous equations that in principle can be solved to deduce all state values. A key caveat for this dynamic programming approach to work is that one must have complete knowledge of the environment dynamics $p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)$ (similar remarks apply to the qualities of different state-action pairs).

Problem: The above Bellman equations apply to a generic agent policy $\pi$; how do they specialize to the case of an optimal policy $\pi^*$?

Solution: Reverting back to the form of the “pre-Bellman” equations and using the key insight that the optimal policy $\pi^*(\mathbf f_t|\mathbf x_t)=\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q_{\pi^*}(\mathbf x_t,\mathbf f_t)}$ should only assign non-zero probability to actions of the best quality, one has:

\[v_{\pi^*}(\mathbf x_t)=\text{max}_{\mathbf f_t}q_{\pi^*}(\mathbf x_t,\mathbf f_t)\]

\[q_{\pi^*}(\mathbf x_t,\mathbf f_t)=\sum_{r_{t+1},\mathbf x_{t+1}}p(r_{t+1},\mathbf x_{t+1}|\mathbf x_t,\mathbf f_t)(r_{t+1}+\gamma v_{\pi^*}(\mathbf x_{t+1}))\]

which lead to the Bellman optimality equations.

Problem: In light of Bellman optimality, discuss how policy evaluation and policy improvement work. Hence, explain the notion of generalized policy iteration (GPI).

Solution: Imagine the space of all $(\pi,v)$ pairs, where $\pi(\mathbf f_t|\mathbf x_t)$ is any policy distribution and $v(\mathbf x_t)$ is the value function of some policy. Within this space, there are $2$ canonical subspaces. One is the set $(\pi,v_{\pi})$ of all pairs where $v_{\pi}$ is specifically the value function associated to policy $\pi$, and the other subspace $(\pi_v,v)$ where $\pi_v(\mathbf f_t|\mathbf x_t):=\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q(\mathbf x_t,\mathbf f_t)}$ is greedy with respect to the value function $v$, where $q(\mathbf x_t,\mathbf f_t)=\langle r_{t+1}+\gamma v(\mathbf x_{t+1})|\mathbf x_t,\mathbf f_t\rangle$ (in particular, notice $q\neq q_{\pi}$; the policy $\pi$ does not appear in the definition of $q$).

Policy evaluation is an algorithm for projecting $(v,\pi)\mapsto (v_{\pi},\pi)$. The idea is to view $v$ as a random initial guess for the underlying true policy value function $v_{\pi}$ (though any terminal states $\mathbf x_T$ should be initialized $v(\mathbf x_T):=0$). Then, sweeping through each state $\mathbf x_t$, one updates the value of $v(\mathbf x_t)$ using the Bellman equation involving the (known) randomly initialized values of the other states:

\[v(\mathbf x_t)\mapsto\sum_{\mathbf f_t}\pi(\mathbf f_t|\mathbf x_t)\sum_{\mathbf x_{t+1},r_{t+1}}p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)(r+\gamma v(\mathbf x_{t+1}))\]

With sufficiently many sweeps over the state space, general theorems guarantee convergence $v\to v_{\pi}$.

Policy improvement is an algorithm for projecting onto the other canonical subspace $(v,\pi)\mapsto (v,\pi_v)$. Essentially just act greedy:
\[\pi(\mathbf f_t|\mathbf x_t)\mapsto\delta_{\mathbf f_t,\text{argmax}_{\mathbf f_t}q_{\pi}(\mathbf x_t,\mathbf f_t)}\]
and is rigorously justified by the policy improvement theorem.

Policy iteration roughly amounts to alternating between the $2$ steps of policy evaluation and policy improvement; however there is some leeway in how one does this, for instance one need not necessarily run the policy evaluation step to completion but rather just perform a single sweep over the state space each time (this more “generalized” version of policy iteration is called value iteration, and can be shown to eventually still funnel/converge onto the optimal policy $\pi^*$ as one can prove that $\pi$ is greedy with respect to its own value function $v_{\pi}$ iff $\pi=\pi^*$).

Problem: Distinguish between model-free vs. model-based methods/algorithms in reinforcement learning.

Solution: In both cases, the fundamental limitation (which previously was taken for granted) is incomplete knowledge of the environment dynamics $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$; instead, one can only sample trajectories $\mathbf x_{t=0},\mathbf f_{t=0},r_{t=0},\mathbf x_{t=1},\mathbf f_{t=1},r_{t=1},…$ when running a policy $\pi$ through the MDP. Recalling the fundamental objective of RL is to compute $\pi^*=\text{argmax}_{\pi}\langle R_t|\pi\rangle$, a model-free method for doing so is one which does not attempt to estimate $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$ (here $p$ is what the word “model” refers to, i.e. a model of the environment dynamics) whereas model-based methods do attempt to estimate $p(\mathbf x_{t+1},r_{t+1}|\mathbf x_t,\mathbf f_t)$ and thereby obtain $\pi^*$ through GPI as outlined earlier (see the OpenAI Spinning Up documentation for a nice diagram of this taxonomy and further comments).

Problem: Explain how Monte Carlo evaluation and Monte Carlo control are both model-free methods that accomplish the same actions as their respective model-based cousins, i.e. policy evaluation and policy improvement.

Solution: Fix $\pi$. The key is to recognize that $v_{\pi}(\mathbf x_t)$ is an expectation of the return random variable starting at $\mathbf x_t$, and Monte Carlo point estimators can always be applied to approximate expectation values:

\[v_{\pi}(\mathbf x_t)=\langle R_t|\mathbf x_t,\pi\rangle\approx\frac{1}{N_{\mathbf x_t}}\sum_{n=1}^{N_{\mathbf x_t}}R_n\]

where $N_{\mathbf x_t}$ is the number of times the MDP passes through state $\mathbf x_t$ in however many trajectories, and $R_n$ is the return following the $n^{\text{th}}$ pass through $\mathbf x_t$ within some MDP trajectory.

Lemma: the arithmetic mean $\hat{\mu}_n:=\frac{1}{n}\sum_{i=1}^nx_i$ obeys the recurrence relation $\hat{\mu}_n=\hat{\mu}_{n-1}+\frac{1}{n}(x_n-\hat{\mu}_{n-1})$

Proof: trivial, but instructive to highlight the convex linear combination nature of the recurrence $\hat{\mu}_n=\left(1-\frac{1}{n}\right)\hat{\mu}_{n-1}+\frac{1}{n}x_n$, so for large $n\gg 1$, the majority $1-1/n$ of the weight is given to the prior value of the mean $\hat{\mu}_{n-1}$, only a tiny sliver $1/n$ is given to the new sample $x_n$.

Thus, rather than waiting for a bunch of trajectories to finish before applying the Monte Carlo estimate above, one can introduce an intermediate estimator initialized at $\hat v_{\pi}(\mathbf x_t):=0$ that applies (at the termination of each trajectory) the following iterative rule for each time $n=1,2,…,N_{\mathbf x_t}$ that $\mathbf x_t$ was visited:

\[\hat v_{\pi}(\mathbf x_t)\mapsto\hat v_{\pi}(\mathbf x_t)+\alpha_n(R_n-\hat v_{\pi}(\mathbf x_t))\]

where $\alpha_n:=1/n$. Since this is just a mathematical decomposition of the original Monte Carlo estimate, it is an equally valid model-free method for $v_{\pi}$ estimation. Now comes a curveball: one common modification is to replace $\alpha_n\mapsto\alpha$ with an $n$-independent constant $\alpha\in[0,1]$ (this is called constant-$\alpha$ MC).

If one were presented with the recurrence relation $\hat{\mu}_n=\hat{\mu}_{n-1}+\frac{1}{n}(x_n-\hat{\mu}_{n-1})$ and initial condition $\hat{\mu}_0=0$, one could claim the explicit solution to be $\hat{\mu}_n=\frac{1}{n}\sum_{i=1}^nx_i$. However, if instead one were presented with the recurrence relation $\hat{\mu}’_n=\hat{\mu}’_{n-1}+\alpha(x_n-\hat{\mu}’_{n-1})$ and same initial condition $\hat{\mu}’_0=0$, what would be the solution for $\hat{\mu}’_n$? It’s easy to check $\hat{\mu}’_n=\alpha\sum_{i=1}^n(1-\alpha)^{n-i}x_i$; thus, recent samples are given exponentially more weight, and it’s biased towards underestimating $\langle\hat{\mu}’_n\rangle-\mu=-\mu(1-\alpha)^n$, though is asymptotically unbiased $\lim_{n\to\infty}\langle\hat{\mu}’_n\rangle=\mu$. However, unlike the asymptotically vanishing variance $\sigma^2_{\hat{\mu}_n}=\sigma^2/n$ of the standard mean estimator $\hat{\mu}_n$, the variance $\sigma^2_{\hat{\mu}’_n}=\frac{\alpha\sigma^2}{2-\alpha}(1-(1-\alpha)^{2n})\to\frac{\alpha\sigma^2}{2-\alpha}\approx\alpha\sigma^2/2$ of $\hat{\mu}’_n$ is bounded from below, hence the estimate $\hat{\mu}’_n$ will forever oscillate about $\mu$. It should make sense $\sigma^2_{\hat{\mu}’_n}\propto\alpha$, i.e. a larger step size $\alpha$ means bigger oscillations, though it also means faster convergence of $\hat{\mu}’_n\to\mu$ because of the $(1-\alpha)^n$.

This all begs the question: why bother with constant-$\alpha$ MC? (come back to this…)

Having used constant-$\alpha$ MC to estimate $v_{\pi}(\mathbf x_t)$ in a model-free way for some fixed agent policy $\pi$, one can naturally ask about model-free estimation of $q_{\pi}(\mathbf x_t,\mathbf f_t)$. It turns out to just be:

\[\hat q_{\pi}(\mathbf x_t,\mathbf f_t)\mapsto\hat q_{\pi}(\mathbf x_t,\mathbf f_t)+\alpha(R_n-\hat q_{\pi}(\mathbf x_t,\mathbf f_t))\]

Conceptually, the standard way to convince oneself of this is that because $\pi$ was fixed, it could be thought of as part of the incomplete knowledge of the environment dynamics $p$ so that the agent never takes any actions but just lets the environment push its state around. This perspective shift from MDP trajectories $\mathbf x_{t=0},\mathbf f_{t=0},r_{t=0},\mathbf x_{t=1},\mathbf f_{t=1},r_{t=1},…$ to MRP trajectories $\mathbf x_{t=0},r_{t=0},\mathbf x_{t=1},r_{t=1},…$ amounts to viewing the MRP states $\mathbf x_t$ as corresponding to augmented state-action pairs $(\mathbf x_t,\mathbf f_t)$ in the MDP. But then it is mathematically apparent that $v_{\pi}(\mathbf x_t)$ in the MRP is just $q_{\pi}(\mathbf x_t,\mathbf f_t)$ in the MDP, so constant-$\alpha$ MC works for $q_{\pi}$-evaluation! Thus, Monte Carlo control chooses to be greedy with respect to $q_{\pi}$, and since $q_{\pi}$ could be estimated in a model-free way, this therefore implements model-free policy improvement.

Problem: Explain the exploration-exploitation tradeoff in reinforcement learning and how a soft policy $\pi$ can help to address this.

Solution: Being greedily exploiting high value states all the time could lead to the agent settling into a “locally” optimal but not globally optimal policy $\pi^*$; the only way to know for sure is to explore all state-action sequences, for instance by ensuring that one’s policy is soft $\pi(\mathbf f_t|\mathbf x_t)>0$ for all states $\mathbf x_t$ and actions $\mathbf f_t$. A crude way to guarantee softness during Monte Carlo control is to use an $\varepsilon$-greedy policy (a better name would be either a $(1-\varepsilon)$-greedy policy or an $\varepsilon$-soft policy), though there are other more refined ways to do so.

Problem: Having distinguished between model-based and model-free methods in RL, one can make an orthogonal distinction between on-policy and off-policy methods. Explain and motivate this taxonomy.

Solution: The constant-$\alpha$ MC $q_{\pi}$-evaluation followed by $\varepsilon$-greedy MC control is a simple model-free, on-policy RL algorithm for discovering $\pi^*$. This is because the trajectory samples used to evaluate $q_{\pi}$ came from the same policy $\pi$ that is subsequently being improved. However, in practice the trajectory samples may come from some (known) behavior policy $\pi_b$ even though the goal is to evaluate $q_{\pi}$ so as to improve another, different (known) target policy $\pi\neq\pi_b$. This is where off-policy methods come in (cf. the distinction between on-shell vs. off-shell particles in physics). The obvious way to bridge $\pi_b$ and $\pi$ is via importance sampling:

\[q_{\pi}(\mathbf x_t,\mathbf f_t)=\langle R_t|\mathbf x_t,\mathbf f_t,\pi\rangle\]

\[=\frac{1}{N_{\mathbf x_t}}\sum_{n=1}^{N_{\mathbf x_t}}\frac{\pi(\mathbf f_{t+1}|\mathbf x_{t+1})…\pi(\mathbf f_T|\mathbf x_T)}{\pi_b(\mathbf f_{t+1}|\mathbf x_{t+1})…\pi_b(\mathbf f_T|\mathbf x_T)}R_n\]

where it is essential to emphasize that the trajectories in this sum are sampled under the behavior policy $\pi_b$! This is model-free because all the transition probabilities have cancelled out in the importance sampling quotient.

Posted in Blog | Leave a comment

Density Functional Theory

Posted on January 28, 2026 by wdengquantum.me

Problem: In one sentence, what is the essence of DFT?

Solution: To replace $\Psi\mapsto n$, where the number density of a system of $N$ identical quantum particles (usually electrons) $n(\mathbf x)$ is:

\[n(\mathbf x):=N\int d^3\mathbf x_2…d^3\mathbf x_N |\Psi(\mathbf x,\mathbf x_2,…,\mathbf x_N)|^2\]

in terms of its $N$-body wavefunction $\Psi(\mathbf x,\mathbf x_2,…,\mathbf x_N)$. Thus, $\int d^3\mathbf x n(\mathbf x)=N$ is normalized.

Problem: Write down the non-relativistic Hamiltonian $H$ of a molecule. Attempt to express the expected energy $E=\langle\Psi|H|\Psi\rangle$ in many-body state $|\Psi\rangle\in\bigwedge^NL^2(\mathbf R^3\to\mathbf C)\otimes\mathbf C^2$ in terms of $n(\mathbf x)$ defined above; this functional $E=E[n]$ of the density $n$ explains the name “density functional theory“.

Solution: (the Born-Oppenheimer approximation is implicitly used whereby the electrons move on an effective/external potential energy surface $V_{\text{ext}}(\mathbf x)$ defined by a collection of stationary nuclei)

\[H=\sum_{i=1}^N\frac{|\mathbf P_i|^2}{2m}+V_{\text{ext}}(\mathbf X_i)+\frac{1}{2}\sum_{1\leq i\neq j\leq N}\frac{\alpha\hbar c}{|\mathbf X_i-\mathbf X_j|}\]

One can check that:

\[E=\frac{\hbar^2}{2m}\sum_{i=1}^N\int d^3\mathbf x_1…d^3\mathbf x_N\biggr|\frac{\partial\Psi}{\partial\mathbf x_i}\biggr|^2+\int d^3\mathbf x n(\mathbf x)V_{\text{ext}}(\mathbf x)+\frac{\alpha\hbar c}{2}\int d^3\mathbf x d^3\mathbf x’\frac{n_2(\mathbf x,\mathbf x’)}{|\mathbf x-\mathbf x’|}\]

where $n_2(\mathbf x,\mathbf x’)=N(N-1)\int d^3\mathbf x_3…d^3\mathbf x_N|\Psi(\mathbf x,\mathbf x’,\mathbf x_3,…,\mathbf x_N)|^2$ so in particular $(N-1)n(\mathbf x)=\int d^3\mathbf x’ n_2(\mathbf x,\mathbf x’)$.

Problem: Write down the Thomas-Fermi approximation for $E[n]$, and explain why it sucks. Also derive the von Weizsäcker correction to the Thomas-Fermi energy functional.

Solution: The Thomas-Fermi functional treats the electrons as a locally ideal Fermi gas, thus having the usual Pauli pressure $p\propto n^{5/3}$ in $\mathbf R^3$ and thus (kinetic) energy density $3p/2$ scaling likewise. Furthermore, the exchange-correlation functional is taken to vanish $E_{\text{xc}}[n]=0$ (thus, Thomas-Fermi is often said to be semiclassical):

\[E_{\text{TF}}[n]=\frac{3\hbar^2(3\pi^2)^{2/3}}{10m}\int d^3\mathbf x n^{5/3}(\mathbf x)+\int d^3\mathbf x n(\mathbf x)V_{\text{ext}}(\mathbf x)+\frac{\alpha\hbar c}{2}\int d^3\mathbf x d^3\mathbf x’\frac{n(\mathbf x)n(\mathbf x’)}{|\mathbf x-\mathbf x’|}\]

One has:

\[\frac{\delta}{\delta n(\mathbf x)}\left(E_{\text{TF}}-\mu\left(\int d^3\mathbf xn(\mathbf x)-N\right)\right)=\frac{\hbar^2k^2_F(\mathbf x)}{2m}+V_{\text{eff}}(\mathbf x)-\mu=0\]

where $k_F(\mathbf x):=(3\pi^2n(\mathbf x))^{1/3}$ and $V_{\text{eff}}(\mathbf x)=V_{\text{ext}}(\mathbf x)+\alpha\hbar c\int d^3\mathbf x’\frac{n(\mathbf x’)}{|\mathbf x-\mathbf x’|}$.

(aside: although the Thomas-Fermi model is meant to be applied to molecules, a result of Teller showed that $E_{\text{TF molecule}}>2E_{\text{TF atom}}$ so Thomas-Fermi theory predicts the molecular state is unstable with respect to dissociation. In a central external potential such as due to a single nucleus $V_{\text{ext}}(r)=-Z\alpha\hbar c/r$, the effective potential $V_{\text{eff}}(r)$ and electron density $n(r)$ are both isotropic as well. Starting from Poisson’s equation:

\[\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2\frac{V_{\text{eff}}(r)}{-e}=-\frac{-en(r)}{\varepsilon_0}\]

with $\biggr|\frac{\partial}{\partial\mathbf x}\biggr|^2f(r)=\frac{1}{r^2}\frac{d}{dr}\left(r^2\frac{df}{dr}\right)=\frac{1}{r}\frac{d^2}{dr^2}(rf(r))$ (use the latter) one can check that by defining $f(r):=\left(\frac{8}{3\pi a_0}\right)^2rk^2_F(r)$, one recovers the Thomas-Fermi ODE:

\[\frac{d^2 f}{dr^2}=\frac{f^{3/2}(r)}{\sqrt{r}}\]

The von Weizsäcker term is a correction to the kinetic energy functional in the Thomas-Fermi model that accounts for density gradients:

\[T_{\text{vW}}[n]=\frac{\hbar^2}{8m}\int d^3\mathbf x\frac{|\partial n/\partial\mathbf x|^2}{n(\mathbf x)}\]

Heuristically, it can be justified by writing the expected kinetic energy of a single electron $\frac{\hbar^2}{2m}\int d^3\mathbf x |\partial\psi/\partial\mathbf x|^2$ with wavefunction $\psi(\mathbf x)\in\mathbf R$ and replacing $\psi(\mathbf x)\mapsto\sqrt{n(\mathbf x)}$.

Problem: Prove the two Hohenberg-Kohn theorems.

Solution: First, it is useful to prove:

Lemma: If $V_{\text{ext}}(\mathbf x)\neq V’_{\text{ext}}(\mathbf x)$ represent $2$ physically distinct external potentials (i.e. they differ by more than just some additive constant), then the corresponding molecular Hamiltonians $H\neq H’$ do not share any non-zero eigenstates.
Proof: Suppose for sake of contradiction that there exists $|\Psi\rangle\neq 0$ such that
\[H|\Psi\rangle=E|\Psi\rangle\]
\[H’||\Psi\rangle=E’|\Psi\rangle\]
Subtracting, one obtains:
\[(H’-H)|\Psi\rangle=(E’-E)|\Psi\rangle\]
Inspecting the form of the molecular Hamiltonian $H$ given earlier, one sees that the “universal” part (i.e. the kinetic energy and electron-electron Coulomb repulsion terms) cancel, leaving only the non-universal residue $H’-H=\sum_{i=1}^NV’_{\text{ext}}(\mathbf X_i)-V_{\text{ext}}(\mathbf X_i)$. But the RHS $E’-E$ is $\mathbf X_i$-independent for all $i=1,…,N$; this means $V’_{\text{ext}}(\mathbf x)-V_{\text{ext}}(\mathbf x)=\text{const}$ which is a contradiction.

With this in mind, the first HK theorem follows by specializing to the respective (distinct by the above lemma) ground states $|\Psi_0\rangle\neq|\Psi’_0\rangle$ of $H$ and $H’$. Apply the variational bound both ways:

\[\]

Problem: Demonstrate the first Hohenberg-Kohn theorem for $2$ electrons in a $1$-dimensional harmonic trap.

Solution:

Problem: Derive the Kohn-Sham equations.

Solution: Fictitious ideal electron gas with density $n(\mathbf x)$.

Problem:

Posted in Blog | Leave a comment

PyTorch Fundamentals (Part $2$)

Posted on January 27, 2026 by wdengquantum.me

Problem: Do an end-to-end walkthrough of the PyTorch machine learning workflow using the most basic univariate linear regression example. In particular, generate some linear data over a normalized feature space (whose slope $w$ and intercept $b$ would in practice be a priori unknown), split that linear data into training and testing subsets (no cross-validation dataset needed for this simple example), define the linear layer class, instantiate a model object of the class, and starting from random values of $w$ and $b$, use stochastic gradient descent with learning rate $\alpha=0.01$ to minimize the training cost function $C_{\text{train}}(w,b)$ based on an $L^1$ loss. Iterate SGD for $300$ epochs, and for every $20$ epochs, record the current value of $C_{\text{train}}(w,b)$ and the current value of $C_{\text{test}}(w,b)$. Plot these cost function curves as a function of the epoch number $0, 20, 40,…$. Save the final state dictionary of the model’s learned parameters $w,b$ post-training, and load it back onto a new instance of the model class.

Solution:

pytorch_workflow

Typical PyTorch Workflow

Get data ready (turn into tensors).
Build of pick a pretrained model (to suit one’s problem).
Fit the model to the data and make a prediction.
Evaluate the model
Improve through experimentation.
Save and reload your trained model

In [1]:

import torch
from torch import nn # nn module contains all of PyTorch's building blocks for designing architectures
import matplotlib.pyplot as plt 

torch.__version__

Out[1]:

'2.9.1+cu128'

In [2]:

# Create known parameters
weight = 0.7
bias = 0.3 
start = 0
end = 1
step = 0.02
x = torch.arange(start, end, step).unsqueeze(dim=1)
#print(x)
y = weight * x + bias
print(f"First 10 elements of x: {x[:10]}")
print(f"First 10 elements of y: {y[:10]}")
plt.plot(x, y)

First 10 elements of x: tensor([[0.0000],
        [0.0200],
        [0.0400],
        [0.0600],
        [0.0800],
        [0.1000],
        [0.1200],
        [0.1400],
        [0.1600],
        [0.1800]])
First 10 elements of y: tensor([[0.3000],
        [0.3140],
        [0.3280],
        [0.3420],
        [0.3560],
        [0.3700],
        [0.3840],
        [0.3980],
        [0.4120],
        [0.4260]])

Out[2]:

[<matplotlib.lines.Line2D at 0x7f60957a38b0>]

In [3]:

train_split = int(0.8 * len(x)) # or int(0.8 * len(y))
x_train, y_train = x[:train_split], y[:train_split]
x_test, y_test = x[train_split:], y[train_split:]
# can also use scikit learn's splitting method which 
# adds some randomness to the training data.
print(len(x_train), len(y_train), len(x_test), len(y_test))

40 40 10 10

In [4]:

# Create linear regression model class
class LinearRegressionModel(nn.Module):
    #almost all classes inherit from nn.Module
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(1,
                                                requires_grad=True, #requires_grad=True is default
                                                dtype=torch.float))
        self.bias = nn.Parameter(torch.randn(1,
                                             requires_grad=True,
                                             dtype=torch.float))
        # or, instead of hard coding the weights and biases, PyTorch's nn.Module class has
        # a built-in nn.Linear layer, so above can be replaced by something like
        # self.linear_layer = nn.Linear(in_features=1, out_features=1)
        # Forward method to define computation in the model, should always be defined to override 
        # default method in nn.Module
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.weights * x + self.bias

PyTorch model building essentials¶

torch.nn – contains all the building blocks for computational graphs (e.g. neural networks)
torch.nn.Parameter – what parameters should our model learn, often a PyTorch layer from torch.nn will set these for us
torch.nn.Module – The base class for all neural network modules, if you subclass/inherit from it, you should overwrite forward()
torch.optim – this where the optimizers in PyTorch live, they will help with gradient descent
def forward() – All nn.Module subclasses require one to overwrite this method; this defines what happens in forward pass/propagation/computation.

In [5]:

# Checking contents of PyTorch model
torch.manual_seed(42)
# Create object instance of the class:
model_0 = LinearRegressionModel()
print(list(model_0.parameters()))
print(model_0.state_dict())

with torch.inference_mode(): #inference mode basically tells the model that this is now training/cross-validation data, not to be used for updating parameters
    y_hat = model_0(x_test)

plt.scatter(x, y)
plt.scatter(x_test, y_hat)

[Parameter containing:
tensor([0.3367], requires_grad=True), Parameter containing:
tensor([0.1288], requires_grad=True)]
OrderedDict([('weights', tensor([0.3367])), ('bias', tensor([0.1288]))])

Out[5]:

<matplotlib.collections.PathCollection at 0x7f60936cb610>

In [6]:

# Setup loss function
loss_fn = nn.L1Loss() #mean absolute deviation
print(loss_fn)
# Setup optimizer
optimizer = torch.optim.SGD(params=model_0.parameters(),
                            lr=0.01)

L1Loss()

In [7]:

N_epochs = 200 #hyperparameter

epoch_count = []
train_loss_values = []
test_loss_values = []

# Training loop
for epoch in range(N_epochs):
    model_0.train() # put model in training mode (default state)

    # Forward pass
    y_hat = model_0(x_train)

    # Calculate loss
    loss = loss_fn(y_hat, y_train)
    print(loss)
    # Optimizer zero grad
    optimizer.zero_grad()

    # Backpropagation
    loss.backward()

    # Step the optimizer (gradient descent)
    optimizer.step() # by default optimizer changes will accumulate through the loop, so need to zero the gradient in above step

    # Testing loop
    model_0.eval()
    with torch.inference_mode():
        test_pred = model_0(x_test)
        test_loss = loss_fn(test_pred, y_test)

    if epoch % 10 == 0:
        epoch_count.append(epoch)
        train_loss_values.append(loss)
        test_loss_values.append(test_loss)
        print(f"Epoch: {epoch} | Loss: {loss} | Test Loss: {test_loss}")

    print(model_0.state_dict())

tensor(0.3129, grad_fn=<MeanBackward0>)
Epoch: 0 | Loss: 0.31288138031959534 | Test Loss: 0.48106518387794495
OrderedDict([('weights', tensor([0.3406])), ('bias', tensor([0.1388]))])
tensor(0.3014, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3445])), ('bias', tensor([0.1488]))])
tensor(0.2898, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3484])), ('bias', tensor([0.1588]))])
tensor(0.2783, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3523])), ('bias', tensor([0.1688]))])
tensor(0.2668, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3562])), ('bias', tensor([0.1788]))])
tensor(0.2553, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3601])), ('bias', tensor([0.1888]))])
tensor(0.2438, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3640])), ('bias', tensor([0.1988]))])
tensor(0.2322, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3679])), ('bias', tensor([0.2088]))])
tensor(0.2207, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3718])), ('bias', tensor([0.2188]))])
tensor(0.2092, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3757])), ('bias', tensor([0.2288]))])
tensor(0.1977, grad_fn=<MeanBackward0>)
Epoch: 10 | Loss: 0.1976713240146637 | Test Loss: 0.3463551998138428
OrderedDict([('weights', tensor([0.3796])), ('bias', tensor([0.2388]))])
tensor(0.1862, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3835])), ('bias', tensor([0.2488]))])
tensor(0.1746, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3874])), ('bias', tensor([0.2588]))])
tensor(0.1631, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3913])), ('bias', tensor([0.2688]))])
tensor(0.1516, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3952])), ('bias', tensor([0.2788]))])
tensor(0.1401, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.3991])), ('bias', tensor([0.2888]))])
tensor(0.1285, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4030])), ('bias', tensor([0.2988]))])
tensor(0.1170, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4069])), ('bias', tensor([0.3088]))])
tensor(0.1061, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4108])), ('bias', tensor([0.3178]))])
tensor(0.0968, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4146])), ('bias', tensor([0.3258]))])
tensor(0.0891, grad_fn=<MeanBackward0>)
Epoch: 20 | Loss: 0.08908725529909134 | Test Loss: 0.21729660034179688
OrderedDict([('weights', tensor([0.4184])), ('bias', tensor([0.3333]))])
tensor(0.0823, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4222])), ('bias', tensor([0.3403]))])
tensor(0.0764, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4258])), ('bias', tensor([0.3463]))])
tensor(0.0716, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4293])), ('bias', tensor([0.3518]))])
tensor(0.0675, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4328])), ('bias', tensor([0.3568]))])
tensor(0.0640, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4361])), ('bias', tensor([0.3613]))])
tensor(0.0610, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4394])), ('bias', tensor([0.3653]))])
tensor(0.0585, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4425])), ('bias', tensor([0.3688]))])
tensor(0.0564, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4455])), ('bias', tensor([0.3718]))])
tensor(0.0546, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4483])), ('bias', tensor([0.3743]))])
tensor(0.0531, grad_fn=<MeanBackward0>)
Epoch: 30 | Loss: 0.053148526698350906 | Test Loss: 0.14464017748832703
OrderedDict([('weights', tensor([0.4512])), ('bias', tensor([0.3768]))])
tensor(0.0518, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4539])), ('bias', tensor([0.3788]))])
tensor(0.0507, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4564])), ('bias', tensor([0.3803]))])
tensor(0.0498, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4590])), ('bias', tensor([0.3818]))])
tensor(0.0490, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4615])), ('bias', tensor([0.3833]))])
tensor(0.0482, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4639])), ('bias', tensor([0.3843]))])
tensor(0.0475, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4662])), ('bias', tensor([0.3853]))])
tensor(0.0469, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4684])), ('bias', tensor([0.3858]))])
tensor(0.0464, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4706])), ('bias', tensor([0.3863]))])
tensor(0.0459, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4728])), ('bias', tensor([0.3868]))])
tensor(0.0454, grad_fn=<MeanBackward0>)
Epoch: 40 | Loss: 0.04543796554207802 | Test Loss: 0.11360953003168106
OrderedDict([('weights', tensor([0.4748])), ('bias', tensor([0.3868]))])
tensor(0.0450, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4768])), ('bias', tensor([0.3868]))])
tensor(0.0446, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4788])), ('bias', tensor([0.3868]))])
tensor(0.0442, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4808])), ('bias', tensor([0.3868]))])
tensor(0.0438, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4828])), ('bias', tensor([0.3868]))])
tensor(0.0434, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4848])), ('bias', tensor([0.3868]))])
tensor(0.0431, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4866])), ('bias', tensor([0.3863]))])
tensor(0.0427, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4884])), ('bias', tensor([0.3858]))])
tensor(0.0424, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4902])), ('bias', tensor([0.3853]))])
tensor(0.0420, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4920])), ('bias', tensor([0.3848]))])
tensor(0.0417, grad_fn=<MeanBackward0>)
Epoch: 50 | Loss: 0.04167863354086876 | Test Loss: 0.09919948130846024
OrderedDict([('weights', tensor([0.4938])), ('bias', tensor([0.3843]))])
tensor(0.0413, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4956])), ('bias', tensor([0.3838]))])
tensor(0.0410, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4974])), ('bias', tensor([0.3833]))])
tensor(0.0406, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.4992])), ('bias', tensor([0.3828]))])
tensor(0.0403, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5010])), ('bias', tensor([0.3823]))])
tensor(0.0399, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5028])), ('bias', tensor([0.3818]))])
tensor(0.0396, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5046])), ('bias', tensor([0.3813]))])
tensor(0.0392, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5064])), ('bias', tensor([0.3808]))])
tensor(0.0389, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5082])), ('bias', tensor([0.3803]))])
tensor(0.0385, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5100])), ('bias', tensor([0.3798]))])
tensor(0.0382, grad_fn=<MeanBackward0>)
Epoch: 60 | Loss: 0.03818932920694351 | Test Loss: 0.08886633068323135
OrderedDict([('weights', tensor([0.5116])), ('bias', tensor([0.3788]))])
tensor(0.0379, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5134])), ('bias', tensor([0.3783]))])
tensor(0.0375, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5152])), ('bias', tensor([0.3778]))])
tensor(0.0372, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5168])), ('bias', tensor([0.3768]))])
tensor(0.0368, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5186])), ('bias', tensor([0.3763]))])
tensor(0.0365, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5202])), ('bias', tensor([0.3753]))])
tensor(0.0361, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5220])), ('bias', tensor([0.3748]))])
tensor(0.0358, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5236])), ('bias', tensor([0.3738]))])
tensor(0.0354, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5254])), ('bias', tensor([0.3733]))])
tensor(0.0351, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5272])), ('bias', tensor([0.3728]))])
tensor(0.0348, grad_fn=<MeanBackward0>)
Epoch: 70 | Loss: 0.03476089984178543 | Test Loss: 0.0805937647819519
OrderedDict([('weights', tensor([0.5288])), ('bias', tensor([0.3718]))])
tensor(0.0344, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5306])), ('bias', tensor([0.3713]))])
tensor(0.0341, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5322])), ('bias', tensor([0.3703]))])
tensor(0.0337, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5340])), ('bias', tensor([0.3698]))])
tensor(0.0334, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5355])), ('bias', tensor([0.3688]))])
tensor(0.0330, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5373])), ('bias', tensor([0.3683]))])
tensor(0.0327, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5391])), ('bias', tensor([0.3678]))])
tensor(0.0324, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5407])), ('bias', tensor([0.3668]))])
tensor(0.0320, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5425])), ('bias', tensor([0.3663]))])
tensor(0.0317, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5441])), ('bias', tensor([0.3653]))])
tensor(0.0313, grad_fn=<MeanBackward0>)
Epoch: 80 | Loss: 0.03132382780313492 | Test Loss: 0.07232122868299484
OrderedDict([('weights', tensor([0.5459])), ('bias', tensor([0.3648]))])
tensor(0.0310, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5475])), ('bias', tensor([0.3638]))])
tensor(0.0306, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5493])), ('bias', tensor([0.3633]))])
tensor(0.0303, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5509])), ('bias', tensor([0.3623]))])
tensor(0.0300, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5527])), ('bias', tensor([0.3618]))])
tensor(0.0296, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5545])), ('bias', tensor([0.3613]))])
tensor(0.0293, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5561])), ('bias', tensor([0.3603]))])
tensor(0.0289, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5579])), ('bias', tensor([0.3598]))])
tensor(0.0286, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5595])), ('bias', tensor([0.3588]))])
tensor(0.0282, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5613])), ('bias', tensor([0.3583]))])
tensor(0.0279, grad_fn=<MeanBackward0>)
Epoch: 90 | Loss: 0.02788739837706089 | Test Loss: 0.06473556160926819
OrderedDict([('weights', tensor([0.5629])), ('bias', tensor([0.3573]))])
tensor(0.0275, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5647])), ('bias', tensor([0.3568]))])
tensor(0.0272, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5665])), ('bias', tensor([0.3563]))])
tensor(0.0269, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5681])), ('bias', tensor([0.3553]))])
tensor(0.0265, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5699])), ('bias', tensor([0.3548]))])
tensor(0.0262, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5715])), ('bias', tensor([0.3538]))])
tensor(0.0258, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5733])), ('bias', tensor([0.3533]))])
tensor(0.0255, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5748])), ('bias', tensor([0.3523]))])
tensor(0.0251, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5766])), ('bias', tensor([0.3518]))])
tensor(0.0248, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5784])), ('bias', tensor([0.3513]))])
tensor(0.0245, grad_fn=<MeanBackward0>)
Epoch: 100 | Loss: 0.024458957836031914 | Test Loss: 0.05646304413676262
OrderedDict([('weights', tensor([0.5800])), ('bias', tensor([0.3503]))])
tensor(0.0241, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5818])), ('bias', tensor([0.3498]))])
tensor(0.0238, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5834])), ('bias', tensor([0.3488]))])
tensor(0.0234, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5852])), ('bias', tensor([0.3483]))])
tensor(0.0231, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5868])), ('bias', tensor([0.3473]))])
tensor(0.0227, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5886])), ('bias', tensor([0.3468]))])
tensor(0.0224, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5902])), ('bias', tensor([0.3458]))])
tensor(0.0221, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5920])), ('bias', tensor([0.3453]))])
tensor(0.0217, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5938])), ('bias', tensor([0.3448]))])
tensor(0.0214, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5954])), ('bias', tensor([0.3438]))])
tensor(0.0210, grad_fn=<MeanBackward0>)
Epoch: 110 | Loss: 0.021020207554101944 | Test Loss: 0.04819049686193466
OrderedDict([('weights', tensor([0.5972])), ('bias', tensor([0.3433]))])
tensor(0.0207, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.5988])), ('bias', tensor([0.3423]))])
tensor(0.0203, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6006])), ('bias', tensor([0.3418]))])
tensor(0.0200, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6022])), ('bias', tensor([0.3408]))])
tensor(0.0196, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6040])), ('bias', tensor([0.3403]))])
tensor(0.0193, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6058])), ('bias', tensor([0.3398]))])
tensor(0.0190, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6074])), ('bias', tensor([0.3388]))])
tensor(0.0186, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6092])), ('bias', tensor([0.3383]))])
tensor(0.0183, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6108])), ('bias', tensor([0.3373]))])
tensor(0.0179, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6126])), ('bias', tensor([0.3368]))])
tensor(0.0176, grad_fn=<MeanBackward0>)
Epoch: 120 | Loss: 0.01758546568453312 | Test Loss: 0.04060482233762741
OrderedDict([('weights', tensor([0.6141])), ('bias', tensor([0.3358]))])
tensor(0.0172, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6159])), ('bias', tensor([0.3353]))])
tensor(0.0169, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6175])), ('bias', tensor([0.3343]))])
tensor(0.0166, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6193])), ('bias', tensor([0.3338]))])
tensor(0.0162, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6211])), ('bias', tensor([0.3333]))])
tensor(0.0159, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6227])), ('bias', tensor([0.3323]))])
tensor(0.0155, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6245])), ('bias', tensor([0.3318]))])
tensor(0.0152, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6261])), ('bias', tensor([0.3308]))])
tensor(0.0148, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6279])), ('bias', tensor([0.3303]))])
tensor(0.0145, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6295])), ('bias', tensor([0.3293]))])
tensor(0.0142, grad_fn=<MeanBackward0>)
Epoch: 130 | Loss: 0.014155393466353416 | Test Loss: 0.03233227878808975
OrderedDict([('weights', tensor([0.6313])), ('bias', tensor([0.3288]))])
tensor(0.0138, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6331])), ('bias', tensor([0.3283]))])
tensor(0.0135, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6347])), ('bias', tensor([0.3273]))])
tensor(0.0131, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6365])), ('bias', tensor([0.3268]))])
tensor(0.0128, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6381])), ('bias', tensor([0.3258]))])
tensor(0.0124, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6399])), ('bias', tensor([0.3253]))])
tensor(0.0121, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6415])), ('bias', tensor([0.3243]))])
tensor(0.0118, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6433])), ('bias', tensor([0.3238]))])
tensor(0.0114, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6451])), ('bias', tensor([0.3233]))])
tensor(0.0111, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6467])), ('bias', tensor([0.3223]))])
tensor(0.0107, grad_fn=<MeanBackward0>)
Epoch: 140 | Loss: 0.010716589167714119 | Test Loss: 0.024059748277068138
OrderedDict([('weights', tensor([0.6485])), ('bias', tensor([0.3218]))])
tensor(0.0104, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6501])), ('bias', tensor([0.3208]))])
tensor(0.0100, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6519])), ('bias', tensor([0.3203]))])
tensor(0.0097, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6534])), ('bias', tensor([0.3193]))])
tensor(0.0093, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6552])), ('bias', tensor([0.3188]))])
tensor(0.0090, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6568])), ('bias', tensor([0.3178]))])
tensor(0.0087, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6586])), ('bias', tensor([0.3173]))])
tensor(0.0083, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6604])), ('bias', tensor([0.3168]))])
tensor(0.0080, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6620])), ('bias', tensor([0.3158]))])
tensor(0.0076, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6638])), ('bias', tensor([0.3153]))])
tensor(0.0073, grad_fn=<MeanBackward0>)
Epoch: 150 | Loss: 0.0072835334576666355 | Test Loss: 0.016474086791276932
OrderedDict([('weights', tensor([0.6654])), ('bias', tensor([0.3143]))])
tensor(0.0069, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6672])), ('bias', tensor([0.3138]))])
tensor(0.0066, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6688])), ('bias', tensor([0.3128]))])
tensor(0.0063, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6706])), ('bias', tensor([0.3123]))])
tensor(0.0059, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6724])), ('bias', tensor([0.3118]))])
tensor(0.0056, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6740])), ('bias', tensor([0.3108]))])
tensor(0.0052, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6758])), ('bias', tensor([0.3103]))])
tensor(0.0049, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6774])), ('bias', tensor([0.3093]))])
tensor(0.0045, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6792])), ('bias', tensor([0.3088]))])
tensor(0.0042, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6808])), ('bias', tensor([0.3078]))])
tensor(0.0039, grad_fn=<MeanBackward0>)
Epoch: 160 | Loss: 0.0038517764769494534 | Test Loss: 0.008201557211577892
OrderedDict([('weights', tensor([0.6826])), ('bias', tensor([0.3073]))])
tensor(0.0035, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6844])), ('bias', tensor([0.3068]))])
tensor(0.0032, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6860])), ('bias', tensor([0.3058]))])
tensor(0.0028, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6878])), ('bias', tensor([0.3053]))])
tensor(0.0025, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6894])), ('bias', tensor([0.3043]))])
tensor(0.0021, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6912])), ('bias', tensor([0.3038]))])
tensor(0.0018, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6927])), ('bias', tensor([0.3028]))])
tensor(0.0015, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6947])), ('bias', tensor([0.3028]))])
tensor(0.0012, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
Epoch: 170 | Loss: 0.008932482451200485 | Test Loss: 0.005023092031478882
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
Epoch: 180 | Loss: 0.008932482451200485 | Test Loss: 0.005023092031478882
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
Epoch: 190 | Loss: 0.008932482451200485 | Test Loss: 0.005023092031478882
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])
tensor(0.0089, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6951])), ('bias', tensor([0.2993]))])
tensor(0.0026, grad_fn=<MeanBackward0>)
OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])

In [8]:

with torch.inference_mode(): #inference mode basically tells the model that this is now training/cross-validation data, not to be used for updating parameters
    y_hat = model_0(x_test)

plt.scatter(x, y)
plt.scatter(x_test, y_hat)

Out[8]:

<matplotlib.collections.PathCollection at 0x7f608837beb0>

In [9]:

plt.scatter(epoch_count, torch.tensor(train_loss_values).numpy(), label="Training Cost")
plt.scatter(epoch_count, torch.tensor(test_loss_values).numpy(), label="Testing Cost")
plt.legend()

/tmp/ipykernel_349241/1254886168.py:1: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:836.)
  plt.scatter(epoch_count, torch.tensor(train_loss_values).numpy(), label="Training Cost")

Out[9]:

<matplotlib.legend.Legend at 0x7f60883d6d40>

In [10]:

# Want to be able to save the model parameters! 
# There are 3 main methods for saving + loading model parameters in PyTorch:
print(model_0.state_dict())

from pathlib import Path

# Create a directory to save models into
model_path = Path("models")
model_path.mkdir(parents=True, exist_ok=True)

# Create model save path
model_name = "pytorch_workflow_model_0.pth"
model_save_path = model_path / model_name

# Save the model state_dict
torch.save(model_0.state_dict(), model_save_path)

OrderedDict([('weights', tensor([0.6990])), ('bias', tensor([0.3093]))])

In [11]:

!ls -l models

total 4
-rw-r--r-- 1 william william 2093 Jan 27 13:54 pytorch_workflow_model_0.pth

In [12]:

# Loading a PyTorch model

# Since only saved the model state_dict, will create new
# object instance of the LinearRegressionModel model class and 

loaded_model_0 = LinearRegressionModel()

# Load the saved state_dict of model_0 (this will update the new instance with the updated parameters)

loaded_model_0.load_state_dict(torch.load(f=model_save_path))

Out[12]:

<All keys matched successfully>

In [13]:

next(model_0.parameters()).device

Out[13]:

device(type='cpu')

Exercises!

In [ ]:

#1
import torch
import matplotlib.pyplot as plt

weight = 0.3
bias = 0.9
N = 100
# something I got reminded of in a nasty way is that all features should be normalized!!! When blindly using torch.linspace(0, 100, N), got exploding gradients...
x = torch.linspace(0, 1, N).unsqueeze(dim=1) #the need for unsqueeze is subtle...need for computing cost function C below as y_hat_train and y_hat_test will both be of shape (n, 1) for some n
y = weight * x + bias # in practice don't know how y relates to x
x_train = x[:int(0.8 * N)]
y_train = y[:int(0.8 * N)]
x_test = x[int(0.8 * N):]
y_test = y[int(0.8 * N):]
plt.plot(x_train, y_train, label="Training Data") #hmm...didn't need to convert PyTorch tensors to NumPy arrays for plotting
plt.plot(x_test, y_test, label="Testing Data")
plt.legend()

Out[ ]:

<matplotlib.legend.Legend at 0x7f60a1f5b640>

In [21]:

#2
from torch import nn

torch.manual_seed(42)

class PyTorchModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(1,
                              requires_grad=True,
                              dtype=torch.float))
        self.b = nn.Parameter(torch.randn(1,
                              requires_grad=True,
                              dtype=torch.float))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w * x + self.b
    
model_instance = PyTorchModel()
print(model_instance.state_dict())

#3
C = nn.L1Loss()
optimizer = torch.optim.SGD(params=model_instance.parameters(), lr=0.01)

N_epochs = 300
C_train_curve = []
C_test_curve = []
for i in range(N_epochs):
    # training loop
    model_instance.train()
    y_hat_train = model_instance(x_train) # or model_instance.forward(x_train)
    C_train = C(y_hat_train, y_train)
    optimizer.zero_grad()
    C_train.backward()
    optimizer.step()

    # testing/evaluation loop
    if i % 20 == 0:
        C_train_curve.append(C_train)
        model_instance.eval()
        with torch.inference_mode(): # stop gradient tape recording during testing
            y_hat_test = model_instance(x_test) # or model_instance.forward(x_test)
            C_test = C(y_hat_test, y_test)
            C_test_curve.append(C_test)
        print(f"Epoch: {i}, Training Cost: {C_train}, Test Cost: {C_test}")

print(f"Final Trained Model Parameters: {model_instance.state_dict()}")

OrderedDict([('w', tensor([0.3367])), ('b', tensor([0.1288]))])
Epoch: 0, Training Cost: 0.7565514445304871, Test Cost: 0.7244139909744263
Epoch: 20, Training Cost: 0.524712860584259, Test Cost: 0.45227327942848206
Epoch: 40, Training Cost: 0.29287439584732056, Test Cost: 0.18013274669647217
Epoch: 60, Training Cost: 0.07645779848098755, Test Cost: 0.07569172978401184
Epoch: 80, Training Cost: 0.0533239021897316, Test Cost: 0.11738457530736923
Epoch: 100, Training Cost: 0.046195853501558304, Test Cost: 0.10600712150335312
Epoch: 120, Training Cost: 0.03922543674707413, Test Cost: 0.09009645879268646
Epoch: 140, Training Cost: 0.03225494548678398, Test Cost: 0.07418543100357056
Epoch: 160, Training Cost: 0.02528444491326809, Test Cost: 0.05827441066503525
Epoch: 180, Training Cost: 0.018313953652977943, Test Cost: 0.04236338287591934
Epoch: 200, Training Cost: 0.01134470570832491, Test Cost: 0.025760680437088013
Epoch: 220, Training Cost: 0.004374831914901733, Test Cost: 0.009503781795501709
Epoch: 240, Training Cost: 0.004876463208347559, Test Cost: 0.0061147273518145084
Epoch: 260, Training Cost: 0.004876463208347559, Test Cost: 0.0061147273518145084
Epoch: 280, Training Cost: 0.004876463208347559, Test Cost: 0.0061147273518145084
Final Trained Model Parameters: OrderedDict([('w', tensor([0.3052])), ('b', tensor([0.9028]))])

In [28]:

# 4
import numpy as np
every_20_epochs = np.arange(0, N_epochs, 20)
plt.scatter(every_20_epochs, torch.tensor(C_train_curve).numpy(), label="Training Cost Curve")
plt.scatter(every_20_epochs, torch.tensor(C_test_curve).numpy(), label="Testing Cost Curve")
plt.legend()

#5
from pathlib import Path
model_path = Path("models")
model_path.mkdir(parents=True, exist_ok=True)
model_name = "pytorch_lin_regress_model.pth"
model_save_path = model_path / model_name
torch.save(model_instance.state_dict(), model_save_path)

In [34]:

#6 
another_model_instance = PyTorchModel()
another_model_instance.load_state_dict(torch.load(model_save_path))

another_y_hat = another_model_instance(x)
plt.plot(x, y, label="Original Data (Both Training + Testing)")
plt.plot(x, torch.tensor(another_y_hat).numpy(), label="Loaded model prediction")
plt.legend()

/tmp/ipykernel_349241/188330038.py:7: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
  plt.plot(x, torch.tensor(another_y_hat).numpy(), label="Loaded model prediction")

Out[34]:

<matplotlib.legend.Legend at 0x7f609fe06e90>

In [ ]:

Posted in Blog | Leave a comment

PyTorch Fundamentals (Part $1$)

Posted on January 25, 2026 by wdengquantum.me

Problem: Illustrate some of the basic fundamentals involved in using the PyTorch deep learning library. In particular, discuss the attributes of PyTorch tensors (e.g. dtype, CPU/GPU devices, etc.), how to generate random PyTorch tensors with/without seeding, and operations that can be performed on and between PyTorch tensors.

Solution:

pytorch_fundamentals

In [2]:

import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
print(torch.__version__)

2.9.1+cu128

In [3]:

# scalars are the most basic type of PyTorch tensor
scalar = torch.tensor(8)
print(scalar)
print(scalar.ndim)
print(scalar.item())

tensor(8)
0
8

In [4]:

vector = torch.tensor([7, 7])
print(vector)
print(vector.ndim)
print(vector.shape)

tensor([7, 7])
1
torch.Size([2])

In [5]:

# MATRIX
MATRIX = torch.tensor([[7, 8],
                       [9, 10]])

print(MATRIX)
print(MATRIX.ndim)
print(MATRIX.shape)
print(MATRIX[1])

tensor([[ 7,  8],
        [ 9, 10]])
2
torch.Size([2, 2])
tensor([ 9, 10])

In [6]:

# TENSOR
TENSOR = torch.tensor([[[1, 2, 3], [3, 6, 9], [2, 4, 6]]])
print(TENSOR)
print(TENSOR.ndim)
print(TENSOR.shape)
print(TENSOR.shape[0])

tensor([[[1, 2, 3],
         [3, 6, 9],
         [2, 4, 6]]])
3
torch.Size([1, 3, 3])
1

In [7]:

my_tensor = torch.tensor([[[[[[[[[[[[[[[[[[[[[[[[[[2, 4, 2, 3]]]]]]]]]]]]]]]]]]]]]]]]]])
print(my_tensor.ndim) # number of onion layers of square brackets
print(my_tensor.shape) # number of elements within each onion layer

26
torch.Size([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4])

Random Tensors¶

Why random tensors?

Random tensors are important because the way many neural networks learn is that they start with tensors full of random numbers and then adjust those random numbers to better represent the data.

Start with random numbers -> look at data -> update random numbers -> look at data -> update random numbers

In [8]:

# Creating a random tensor of size (3, 4)
random_tensor = torch.rand(3, 4)
random_tensor

Out[8]:

tensor([[0.0879, 0.6826, 0.4789, 0.1849],
        [0.7974, 0.9331, 0.8372, 0.3934],
        [0.4137, 0.7374, 0.8922, 0.0088]])

In [9]:

random_image_size_tensor = torch.rand(size=(3, 224, 224))
random_image_size_tensor.shape, random_image_size_tensor.ndim

Out[9]:

(torch.Size([3, 224, 224]), 3)

In [10]:

plt.imshow(random_image_size_tensor[1])

Out[10]:

<matplotlib.image.AxesImage at 0x7f2eb56a1930>

In [11]:

zero_tensor = torch.zeros(size=(3, 4))
one_tensor = torch.ones(size=(3, 4))
print(zero_tensor)
print(one_tensor)
print(one_tensor.dtype) # by default, all tensors use float32 (single-point precision) initially
# Multiplication symbol * leads to Hadamard product
print(random_tensor)
print(one_tensor * random_tensor)

tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
torch.float32
tensor([[0.0879, 0.6826, 0.4789, 0.1849],
        [0.7974, 0.9331, 0.8372, 0.3934],
        [0.4137, 0.7374, 0.8922, 0.0088]])
tensor([[0.0879, 0.6826, 0.4789, 0.1849],
        [0.7974, 0.9331, 0.8372, 0.3934],
        [0.4137, 0.7374, 0.8922, 0.0088]])

In [12]:

one_to_ten = torch.arange(start=1, end=11, step=1)
print(one_to_ten)

tensor([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

In [13]:

ten_zeros = torch.zeros_like(input=one_to_ten)
print(ten_zeros)

tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

Tensor datatypes¶

Note: Tensor datatypes is one of the 3 big errors one often runs into in PyTorch:

Tensors not the right datatype
Tensors not the right shape
Tensors not on the right device

In [14]:

# Float 32 tensor
float_32_tensor = torch.tensor([3.0, 6.0, 9.0], 
                               dtype=torch.float32,
                               device=None,
                               requires_grad=False)

print(float_32_tensor)
print(float_32_tensor.dtype)

tensor([3., 6., 9.])
torch.float32

In [15]:

float_16_tensor = float_32_tensor.type(torch.float16)
print(float_16_tensor)
print(float_16_tensor.dtype)

tensor([3., 6., 9.], dtype=torch.float16)
torch.float16

In [16]:

float_16_tensor * float_32_tensor #

Out[16]:

tensor([ 9., 36., 81.])

In [17]:

some_tensor = torch.rand(3, 4)
print(some_tensor)
print(f"Datatype of tensor: {some_tensor.dtype}")
print(f"Shape of tensor: {some_tensor.shape}")
print(f"Device that tensor is on: {some_tensor.device}")

tensor([[0.2600, 0.3511, 0.7676, 0.6426],
        [0.9504, 0.4816, 0.0339, 0.0265],
        [0.9296, 0.9317, 0.4725, 0.8148]])
Datatype of tensor: torch.float32
Shape of tensor: torch.Size([3, 4])
Device that tensor is on: cpu

Manipulating Tensors (Tensor Operation)¶

Addition
Subtraction
Hadamard multiplication
Matrix Multiplication
Division? (Inversion?)

In [72]:

tensor = torch.tensor([1, 2, 3])
print(tensor + 10)
print(tensor * 10)
print(tensor - 10)
print(torch.add(tensor, 10))

tensor([11, 12, 13])
tensor([10, 20, 30])
tensor([-9, -8, -7])
tensor([11, 12, 13])

In [ ]:

tensor1 = torch.tensor([[1, 2],
                        [3, 4]])
tensor2 = torch.tensor([[4, 5],
                        [6, 7]])

print(tensor1 @ tensor2)
print(torch.matmul(tensor1, tensor2))
print(torch.mm(tensor1, tensor2))
%timeit tensor1 @ tensor2

tensor([[16, 19],
        [36, 43]])
tensor([[16, 19],
        [36, 43]])
tensor([[16, 19],
        [36, 43]])
1.88 μs ± 96.6 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

In [ ]:

%%time
# Hard coding matrix multiplication
C = torch.tensor([[0, 0], [0, 0]])
for i in range(2):
    for j in range(2):
        for k in range(2):
            C[i, j] += tensor1[i, k] * tensor2[k, j]

print(C)

tensor([[16, 19],
        [36, 43]])
CPU times: user 1.94 ms, sys: 394 μs, total: 2.34 ms
Wall time: 1.92 ms

Tensor Aggregation

In [ ]:

x = torch.arange(1, 100, 10)
print(x)
print(x.dtype)
print(torch.min(x), x.min())
print(torch.max(x), x.max())
print(torch.mean(x.type(torch.float32)), x.type(torch.float32).mean())
print(torch.sum(x), x.sum())
print(torch.argmin(x), x.argmin())
print(torch.argmax(x), x.argmax())
print(x[9])

tensor([ 1, 11, 21, 31, 41, 51, 61, 71, 81, 91])
torch.int64
tensor(1) tensor(1)
tensor(91) tensor(91)
tensor(46.) tensor(46.)
tensor(460) tensor(460)
tensor(0) tensor(0)
tensor(9) tensor(9)
tensor(91)

In [93]:

# Reshaping, stacking, squeezing, unsqueezing PyTorch tensors
import torch
x = torch.arange(1, 10)
x, x.shape

Out[93]:

(tensor([1, 2, 3, 4, 5, 6, 7, 8, 9]), torch.Size([9]))

In [99]:

x_reshaped = x.reshape(3, 3)
x_reshaped, x_reshaped.shape

Out[99]:

(tensor([[5, 2, 3],
         [4, 5, 6],
         [7, 8, 9]]),
 torch.Size([3, 3]))

In [ ]:

z = x.view(1, 9)
z, z.shape
z[:, 0] = 5
z, x # changes in z will affect x, stored in same memory location

Out[ ]:

(tensor([[5, 2, 3, 4, 5, 6, 7, 8, 9]]), tensor([5, 2, 3, 4, 5, 6, 7, 8, 9]))

In [ ]:

x_stacked = torch.stack([x, x, x, x], dim=1)
x_h_stacked = torch.hstack([x, x, x, x])
x_v_stacked = torch.vstack([x, x, x, x])
print(x_stacked)
print(x_h_stacked)
print(x_v_stacked)

tensor([[5, 5, 5, 5],
        [2, 2, 2, 2],
        [3, 3, 3, 3],
        [4, 4, 4, 4],
        [5, 5, 5, 5],
        [6, 6, 6, 6],
        [7, 7, 7, 7],
        [8, 8, 8, 8],
        [9, 9, 9, 9]])
tensor([5, 2, 3, 4, 5, 6, 7, 8, 9, 5, 2, 3, 4, 5, 6, 7, 8, 9, 5, 2, 3, 4, 5, 6,
        7, 8, 9, 5, 2, 3, 4, 5, 6, 7, 8, 9])
tensor([[5, 2, 3, 4, 5, 6, 7, 8, 9],
        [5, 2, 3, 4, 5, 6, 7, 8, 9],
        [5, 2, 3, 4, 5, 6, 7, 8, 9],
        [5, 2, 3, 4, 5, 6, 7, 8, 9]])
tensor([5, 2, 3, 4, 5, 6, 7, 8, 9])

In [112]:

y = torch.zeros(1, 2, 3, 2, 3)
print(y)
print(y.squeeze(), y.squeeze().shape)

tensor([[[[[0., 0., 0.],
           [0., 0., 0.]],

          [[0., 0., 0.],
           [0., 0., 0.]],

          [[0., 0., 0.],
           [0., 0., 0.]]],


         [[[0., 0., 0.],
           [0., 0., 0.]],

          [[0., 0., 0.],
           [0., 0., 0.]],

          [[0., 0., 0.],
           [0., 0., 0.]]]]])
tensor([[[[0., 0., 0.],
          [0., 0., 0.]],

         [[0., 0., 0.],
          [0., 0., 0.]],

         [[0., 0., 0.],
          [0., 0., 0.]]],


        [[[0., 0., 0.],
          [0., 0., 0.]],

         [[0., 0., 0.],
          [0., 0., 0.]],

         [[0., 0., 0.],
          [0., 0., 0.]]]]) torch.Size([2, 3, 2, 3])

In [124]:

print(x)
print(x.unsqueeze(dim=0))
print(x.unsqueeze(dim=1))

tensor([5, 2, 3, 4, 5, 6, 7, 8, 9])
tensor([[5, 2, 3, 4, 5, 6, 7, 8, 9]])
tensor([[5],
        [2],
        [3],
        [4],
        [5],
        [6],
        [7],
        [8],
        [9]])

In [135]:

random_image = torch.rand(size=(200, 400, 3))
plt.imshow(random_image)
print(f"Shape of random_image: {random_image.shape}")
random_image_permuted = random_image.permute(2, 1, 0)
print(f"Shape of random_image_permuted: {random_image_permuted.shape}")
#random_image_permuted =

Shape of random_image: torch.Size([200, 400, 3])
Shape of random_image_permuted: torch.Size([3, 400, 200])

In [ ]:

# Indexing from PyTorch tensors is similar to indexing from NumPy arrays
x = torch.arange(1, 10).reshape(1, 3, 3)
print(x, x.shape)
print(x[0])
print(x[0][0])
print(x[0][0][0])
print(x[0][:, 2])

tensor([[[1, 2, 3],
         [4, 5, 6],
         [7, 8, 9]]]) torch.Size([1, 3, 3])
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
tensor([1, 2, 3])
tensor(1)
tensor([3, 6, 9])

In [ ]:

# PyTorch tensors & NumPy
np_array = np.arange(1.0, 8.0)
pytorch_tensor = torch.from_numpy(np_array) # since default NumPy dtype is FP64, PyTorch will reflect this in torch.from_numpy()
print(np_array)
print(pytorch_tensor)

T = torch.ones(7)
numpy_arr = T.numpy()
print(T)
print(numpy_arr)
print(numpy_arr.dtype) #now default is FP32 b/c torch tensors default to FP32

[1. 2. 3. 4. 5. 6. 7.]
tensor([1., 2., 3., 4., 5., 6., 7.], dtype=torch.float64)
tensor([1., 1., 1., 1., 1., 1., 1.])
[1. 1. 1. 1. 1. 1. 1.]
float32

In [180]:

random_tensor_A = torch.rand(3, 4)
random_tensor_B = torch.rand(3, 4)

print(random_tensor_A)
print(random_tensor_B)
print(random_tensor_A == random_tensor_B)

tensor([[0.6327, 0.6428, 0.5100, 0.1919],
        [0.9144, 0.7456, 0.6312, 0.8555],
        [0.9164, 0.6492, 0.5402, 0.2073]])
tensor([[0.0675, 0.4836, 0.7717, 0.0868],
        [0.9814, 0.8618, 0.9880, 0.1541],
        [0.8092, 0.0201, 0.0600, 0.7161]])
tensor([[False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]])

In [188]:

# Set a random seed
seed = 42
torch.manual_seed(seed)
random_tensor_C = torch.rand(3, 4)
random_tensor_D = torch.rand(3, 4)
print(random_tensor_C)
print(random_tensor_D)

tensor([[0.8823, 0.9150, 0.3829, 0.9593],
        [0.3904, 0.6009, 0.2566, 0.7936],
        [0.9408, 0.1332, 0.9346, 0.5936]])
tensor([[0.8694, 0.5677, 0.7411, 0.4294],
        [0.8854, 0.5739, 0.2666, 0.6274],
        [0.2696, 0.4414, 0.2969, 0.8317]])

In [190]:

# Running tensors and PyTorch objects on GPUs (and making faster computations)
!nvidia-smi

Sat Jan 24 20:22:27 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.188                Driver Version: 573.71         CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4050 ...    On  |   00000000:01:00.0 Off |                  N/A |
| N/A   25C    P3             14W /   39W |       0MiB /   6141MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

In [ ]:

# Check for GPU access with PyTorch
torch.cuda.is_available()

# If want code to be device-agnostic, then can check if GPU exists:
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
# Count number of GPUs
print(torch.cuda.device_count())

cuda
1

In [ ]:

# Putting tensors & models on GPU
my_tensor = torch.tensor([1, 2, 3])
print(my_tensor, my_tensor.device) #default to CPU

# To move tensor to GPU (if GPU available as encoded in device variable above)
my_tensor_on_gpu = my_tensor.to(device)
print(my_tensor_on_gpu)
print(my_tensor_on_gpu.device)

# Since NumPy only works on CPU, sometimes one may want to perform some arithmetic
# with NumPy arrays, in which case one would like to move a tensor back onto CPU,
# then convert to NumPy array.

my_tensor_on_cpu = my_tensor_on_gpu.to("cpu") # or can do my_tensor_on_gpu.cpu()
print(my_tensor_on_cpu)
print(my_tensor_on_cpu.device)
print(my_tensor_on_cpu.numpy())

tensor([1, 2, 3]) cpu
tensor([1, 2, 3], device='cuda:0')
cuda:0
tensor([1, 2, 3])
cpu
[1 2 3]

Exercise Solutions:

In [374]:

#1 Read the documentation on torch.tensor and torch.Cuda
#2 
random_tensor = torch.rand(size=(7, 7))
print(random_tensor)
#3 
another_random_tensor = torch.rand(size=(1, 7))
print(random_tensor@another_random_tensor.T)

tensor([[0.5159, 0.1636, 0.0958, 0.8985, 0.5814, 0.9148, 0.3324],
        [0.6473, 0.3857, 0.4778, 0.1955, 0.6691, 0.6581, 0.4897],
        [0.3875, 0.1918, 0.8458, 0.1278, 0.7048, 0.3319, 0.2588],
        [0.5898, 0.2403, 0.6152, 0.5982, 0.1288, 0.5832, 0.7130],
        [0.6979, 0.4371, 0.0901, 0.4229, 0.6737, 0.3176, 0.6898],
        [0.8330, 0.2389, 0.5049, 0.7067, 0.5392, 0.5418, 0.5624],
        [0.1069, 0.5393, 0.8462, 0.9506, 0.7939, 0.5670, 0.7335]])
tensor([[1.8984],
        [1.4171],
        [1.1102],
        [1.5038],
        [1.7249],
        [1.8912],
        [2.3290]])

In [ ]:

#4
seed = 1234
torch.manual_seed(seed) #this actually seeds both CPU & GPU with same seed-->same behavior
random_tensor = torch.rand(size=(7, 7), device=device)
#random_tensor = random_tensor.to(device)
print(random_tensor)
another_random_tensor = torch.rand(size=(1, 7), device=device)
#another_random_tensor = another_random_tensor.to(device)
print(random_tensor@another_random_tensor.T)

tensor([[0.1272, 0.8167, 0.5440, 0.6601, 0.2721, 0.9737, 0.3903],
        [0.3394, 0.5451, 0.7312, 0.3864, 0.5959, 0.7578, 0.2126],
        [0.7198, 0.9845, 0.5518, 0.0981, 0.0582, 0.5839, 0.1083],
        [0.9461, 0.3170, 0.8328, 0.6676, 0.2886, 0.9022, 0.8115],
        [0.1784, 0.9534, 0.1486, 0.3882, 0.7977, 0.1752, 0.5777],
        [0.1949, 0.8499, 0.3125, 0.2156, 0.0383, 0.4934, 0.3138],
        [0.3121, 0.5664, 0.1266, 0.7097, 0.0040, 0.5147, 0.2811]],
       device='cuda:0')
tensor([[0.9558],
        [1.2227],
        [0.9335],
        [1.6030],
        [0.9344],
        [0.5282],
        [0.5664]], device='cuda:0')

In [415]:

#4 (seeding the GPU, but then random tensors must be created there!)
seed = 1234
torch.cuda.manual_seed(seed)
random_tensor = torch.rand(size=(7, 7), device=device)
print(random_tensor)
another_random_tensor = torch.rand(size=(1, 7), device=device)
print(random_tensor@another_random_tensor.T)

tensor([[0.1272, 0.8167, 0.5440, 0.6601, 0.2721, 0.9737, 0.3903],
        [0.3394, 0.5451, 0.7312, 0.3864, 0.5959, 0.7578, 0.2126],
        [0.7198, 0.9845, 0.5518, 0.0981, 0.0582, 0.5839, 0.1083],
        [0.9461, 0.3170, 0.8328, 0.6676, 0.2886, 0.9022, 0.8115],
        [0.1784, 0.9534, 0.1486, 0.3882, 0.7977, 0.1752, 0.5777],
        [0.1949, 0.8499, 0.3125, 0.2156, 0.0383, 0.4934, 0.3138],
        [0.3121, 0.5664, 0.1266, 0.7097, 0.0040, 0.5147, 0.2811]],
       device='cuda:0')
tensor([[0.9558],
        [1.2227],
        [0.9335],
        [1.6030],
        [0.9344],
        [0.5282],
        [0.5664]], device='cuda:0')

In [ ]:

#6 (just to prove that torch.manual_seed() seeds both CPU and GPU devices)
torch.manual_seed(1234)
T1_rand = torch.rand(size=(2, 3), device="cuda")
T2_rand = torch.rand(size=(2, 3), device="cuda")
print(T1_rand)
print(T2_rand)
#7
T = T1_rand.T@T2_rand
print(T)
#8
print(torch.max(T), T.max())
print(torch.min(T), T.min())
#9 
print(torch.argmax(T), T.argmax())
print(torch.argmin(T), T.argmin())
#10 
T_rand = torch.rand(size=(1, 1, 1, 10))
print(T_rand, T_rand.shape)
T_rand = T_rand.squeeze()
print(T_rand, T_rand.shape)

tensor([[0.1272, 0.8167, 0.5440],
        [0.6601, 0.2721, 0.9737]], device='cuda:0')
tensor([[0.6208, 0.0276, 0.3255],
        [0.1114, 0.6812, 0.3608]], device='cuda:0')
tensor([[0.1525, 0.4531, 0.2796],
        [0.5374, 0.2079, 0.3640],
        [0.4462, 0.6783, 0.5284]], device='cuda:0')
tensor(0.6783, device='cuda:0') tensor(0.6783, device='cuda:0')
tensor(0.1525, device='cuda:0') tensor(0.1525, device='cuda:0')
tensor(7, device='cuda:0') tensor(7, device='cuda:0')
tensor(0, device='cuda:0') tensor(0, device='cuda:0')
tensor([[[[0.0290, 0.4019, 0.2598, 0.3666, 0.0583, 0.7006, 0.0518, 0.4681,
           0.6738, 0.3315]]]]) torch.Size([1, 1, 1, 10])
tensor([0.0290, 0.4019, 0.2598, 0.3666, 0.0583, 0.7006, 0.0518, 0.4681, 0.6738,
        0.3315]) torch.Size([10])

Posted in Blog | Leave a comment

Self-Attention in Transformers

Posted on December 31, 2025 by wdengquantum.me

Problem: Explain how the transformer architecture works at a mathematical level (e.g. as outlined in the Attention Is All You Need paper).

Solution:

(Tokenization) Partition the inputted natural language text into a sequence of tokens $\tau_1,…,\tau_N$ (here $N$ is bounded from above by the LLMs context size).
(Embedding) Each token $\tau_i$ is embedded as a vector $\tau_i\mapsto\mathbf x_i\in\mathbf R^{n_e}$ that contains information about the token’s generic meaning as well as its position in the inputted natural language text. Here, the hyperparameter $n_e$ is the dimension of the embedding space.
(Single Self-Attention Head) For each embedding vector $\mathbf x_i$, compute its query vector $\mathbf q_i=W_{\mathbf q}\mathbf x_i$, its key vector $\mathbf k_i=W_{\mathbf k}\mathbf x_i$, and its value vector $\mathbf v_i=W_{\mathbf v}\mathbf x_i$. Here, $W_{\mathbf q},W_{\mathbf k}\in\mathbf R^{n_{qk}\times n_e}$ are weight matrices that map from the embedding space $\mathbf R^{n_e}$ to the query/key space $\mathbf R^{n_{qk}}$ of dimension $n_{qk}$ and $W_{\mathbf v}\in\mathbf R^{n_e\times n_e}$ is the weight matrix of values (which in practice is decomposed into a low-rank approximation $W_{\mathbf v}=W_{\mathbf v\uparrow}W_{\mathbf v\downarrow}$ where typically $W_{\mathbf v\downarrow}\in\mathbf R^{n_{qk}\times n_e}$ and $W_{\mathbf v\uparrow}\in\mathbf R^{n_e\times n_{qk}}$). For each $\mathbf x_i$, one computes an update vector $\Delta\mathbf x_i$ to be added to it according to a convex linear combination of the value vectors $\mathbf v_1,…,\mathbf v_N$ of all the embeddings $\mathbf x_1,…,\mathbf x_N$ in the context, specifically:

\[\Delta\mathbf x_i=V\text{softmax}\left(\frac{K^T\mathbf q_i}{\sqrt{n_{qk}}}\right)\]

where $K=(\mathbf k_1,…,\mathbf k_N)\in\mathbf R^{n_{qk}\times N}$ and $V=(\mathbf v_1,…,\mathbf v_N)\in\mathbf R^{n_e\times N}$ are key and value matrices associated to the inputted context (filled with column vectors here rather than the ML convention of row vectors). This map that takes the initial, generic token embeddings $\mathbf x_i$ and nudges them towards more contextualized embeddings $\mathbf x_i\mapsto\mathbf x’_i=\mathbf x_i+\Delta\mathbf x_i$ is called a head of self-attention. The $1/\sqrt{n_{qk}}$ scaling in the softmax temperature is justified on the grounds that if $\mathbf k$ and $\mathbf q$ are random vectors whose independent components each have mean $0$ and variance $1$, then $\mathbf k\cdot\mathbf q$ will have mean $0$ and variance $n_{qk}$, hence the need to normalize by $\sqrt{n_{qk}}$ to ensure $\mathbf k\cdot\mathbf q/\sqrt{n_{qk}}$ continues to have variance $1$.

4. (Multi-Headed Self-Attention) Since context can influence meaning in different ways, repeat the above procedure in parallel for several heads of self-attention; each head will propose a displacement update to each of the $N$ original embeddings $\mathbf x_i$; add up all of them.

5. (Multilayer Perceptron) Linear, ReLU, Linear basically. It is hypothesized that facts are stored in this part of the transformer.

6. (Layers) Alternate between the multi-headed self-attention blocks and MLP blocks, make a probabilistic prediction of the next token $\hat{\tau}_{N+1}$ using only the final, context-rich, modified embedding $\mathbf x’_N$ of the last token $\tau_N$ in the context by applying an unembedding matrix $\mathbf u=W_{\mathbf u}\mathbf x’_N$ and running it through a softmax $\text{softmax}(\mathbf u)$.

Problem: Based on the above discussion of the transformer architecture, explain how a large language model (LLM) like Gemini, ChatGPT, Claude, Grok, DeepSeek, etc. works (at a high level).

Solution: Essentially, since an LLM is a neural network which takes as input some string of text and probabilistically predicts the next token, by seeding it with some corpus of text $T$, the LLM can sample according to the probability distribution it generates for the next token, and append that to $T\mapsto T+\tau$. Then, simply repeat this except pretend that $T+\tau$ was the seed all along. In this way, generative AI models such as ChatGPT (where GPT stands for generative pre-trained transformer) work. In practice, it is helpful to also provide some system prompt like “What follows is a conversation between a user and a knowledgeable AI assistant:”.

Posted in Blog | Leave a comment

JAX Fundamentals (Part $1$)

Posted on December 7, 2025 by wdengquantum.me

JAX_tutorial

$\textbf{Problem}$: What is JAX?

$\textbf{Solution}$: JAX = Autograd + XLA, where Autograd refers to automatic differentiation, and XLA refers to accelerated linear algebra (compiler developed by Google that optimizes code to run fast on GPUs/TPUs).

At a high level, JAX is just NumPy on steroids (and indeed, much of the syntax is identical). At a lower level, JAX is a framework for composable function transformations.

$\textbf{Problem}$: What are the $4$ most important JAX transformations?

$\textbf{Solution}$: jit (just-in-time compilation), grad (gradient), vmap (vectorization), and pmap (parallelization).

The philosophy here is that one can write standard Python functions, and JAX will transform them into GPU/TPU-optimized versions.

$\textbf{Problem}$: Compare JAX vs. PyTorch vs. TensorFlow.

$\textbf{Solution}$: JAX uses a functional programming paradigm whereas PyTorch is object-oriented/imperative and TensorFlow is mixed (Keras is OOP but TF Core is graph). JAX therefore has the steepest learning curve (as it requires unlearning OOP habits). However, it’s XLA performance is arguably the fastest, making it an essential tool for research, etc.

$\textbf{Problem}$: What higher-level libraries built on top of JAX are typically used to build neural networks?

$\textbf{Solution}$: Flax, Haiku, Equinox, Optax, etc.

In [4]:

# JAX's syntax is (for the most part) same as NumPy.
# There is also SciPy API support (jax.scipy)
import jax.numpy as jnp
import numpy as np

# 4 key transform functions
from jax import jit, grad, vmap, pmap

# JAX's low-level API
from jax import lax # anagram of XLA

from jax import make_jaxpr
from jax import random
from jax import device_put

import matplotlib.pyplot as plt

In [5]:

import jax
jax.devices()

Out[5]:

[CudaDevice(id=0)]

In [6]:

# Fact #1: JAX syntax is very similar to NumPy!
x_np = np.linspace(0, 10, 1000)
y_np = 2 * np.sin(x_np) * np.cos(x_np)
plt.plot(x_np, y_np)

Out[6]:

[<matplotlib.lines.Line2D at 0x7f394ad1d180>]

In [7]:

x_jnp = jnp.linspace(0, 10, 1000)
y_jnp = 2 * jnp.sin(x_jnp) * jnp.cos(x_jnp)
plt.plot(x_jnp, y_jnp)

Out[7]:

[<matplotlib.lines.Line2D at 0x7f392ef5d1b0>]

In [8]:

# Fact 2: JAX arrays are immutable! (embrace the functional programming paradigm!)
size = 10
index = 0
value = 23

# NumPy: mutable arrays
x = np.arange(size)
print(x)
x[index] = value
print(x)

[0 1 2 3 4 5 6 7 8 9]
[23  1  2  3  4  5  6  7  8  9]

In [9]:

# JAX: immutable arrays
x = jnp.arange(size)
print(x)
x[index] = value
print(x)

[0 1 2 3 4 5 6 7 8 9]

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[9], line 4
      2 x = jnp.arange(size)
      3 print(x)
----> 4 x[index] = value
      5 print(x)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:596, in _unimplemented_setitem(self, i, x)
    592 def _unimplemented_setitem(self, i, x):
    593   msg = ("JAX arrays are immutable and do not support in-place item assignment."
    594          " Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method:"
    595          " https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html")
--> 596   raise TypeError(msg.format(type(self)))

TypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html

In [ ]:

# Solution:
y = x.at[index].set(value)
print(x)
print(y)

[0 1 2 3 4 5 6 7 8 9]
[23  1  2  3  4  5  6  7  8  9]

In [ ]:

# Fact 3: JAX handles random numbers differently (cf. NumPy)
seed = 0
key = random.PRNGKey(seed)
print(key)
x = random.normal(key,(10,)) # have to pass key here explicitly!
print(type(x),x)

[0 0]
<class 'jaxlib._jax.ArrayImpl'> [ 1.6226422   2.0252647  -0.43359444 -0.07861735  0.1760909  -0.97208923
 -0.49529874  0.4943786   0.6643493  -0.9501635 ]

In [ ]:

# Fact #4: JAX is AI accelerator agnostic (same code runs everywhere!)

size = 3000

# Data is automatically pushed to the AI accelerator.
# No more need for ".to(device)" (PyTorch syntax)
x_jnp = random.normal(key, (size, size), dtype=jnp.float32)
x_np = np.random.normal(size=(size, size)).astype(np.float32) # some diff in API

# block_until_ready() --> ignore time for asynchronous dispatch
%timeit jnp.dot(x_jnp, x_jnp.T).block_until_ready() # on AI accelerator (e.g. GPU) - fast
%timeit np.dot(x_np, x_np.T) # on CPU - slow (NumPy only works with CPUs)
%timeit jnp.dot(x_np, x_np.T).block_until_ready() # on AI accelerator (e.g. GPU) with overhead of transferring np to jnp

x_np_device = device_put(x_np)
%timeit jnp.dot(x_np_device, x_np_device.T).block_until_ready()

7.9 ms ± 90.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
101 ms ± 5.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
22.5 ms ± 1.13 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
7.96 ms ± 39.5 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

$\textbf{Problem}$: What does jit do?

$\textbf{Solution}$: jit is how JAX compiles functions into super-optimized kernels using XLA. It performs a tracing process which makes subsequent calls of the function very fast.

In [ ]:

# Simpler visualizer
def visualize_fn(fn, l=-10, r=10, n=1000):
    x = np.linspace(l, r, num=n)
    y = fn(x)
    plt.plot(x, y); plt.show()

In [ ]:

def selu(x, alpha=1.67, lmbda=1.05): #a type of activation function
    return lmbda * jnp.where(x>0, x, alpha*jnp.exp(x)-alpha)

visualize_fn(selu)

In [ ]:

selu_jit = jit(selu)

data = random.normal(key, (1000000,))
print("Non-jit version")
%timeit selu(data).block_until_ready()
print("Jit version:")
%timeit selu_jit(data).block_until_ready()

Non-jit version
1.66 ms ± 85.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
Jit version:
545 μs ± 117 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

$\textbf{Problem}$: What does grad() do?

$\textbf{Solution}$: Differentiation can be manual/symbolic or numerical. But grad() does it automatically!

In [ ]:

def L(x):
    return jnp.sum(x**3) # this function eats in a vector
# and spits out the sum of cubes of the vector components

x = jnp.arange(0, 6.0, 1)
print(x)
print(L(x)) 
print(grad(L)(x)) #manually, the grad is (3x_1^2,3x_2^2,3x_3^2)

[0. 1. 2. 3. 4. 5.]
225.0
[ 0.  3. 12. 27. 48. 75.]

In [ ]:

# Numeric diff (to check that autodiff works)

def finite_diff(f, x):
    eps = 1e-3
    return jnp.array([(f(x+eps*v)-f(x-eps*v))/(2*eps)
                      for v in jnp.eye(len(x))])

finite_diff(L, x)

Out[ ]:

Array([ 0.       ,  3.0517578, 11.901855 , 26.931763 , 47.98889  ,
       75.07324  ], dtype=float32)

In [ ]:

# Example of autodiff
x = 1.
f = lambda x: x**2 + x + 4
visualize_fn(f, l=-1, r=2, n=100)

In [ ]:

dfdx = grad(f) #2x+1
d2fdx2 = grad(dfdx) #2
d3fdx3 = grad(d2fdx2) #0

print(f(x), dfdx(x), d2fdx2(x), d3fdx3(x))
# More powerful cf. backward() (PyTorch syntax)

6.0 3.0 2.0 0.0

In [ ]:

# Modifying above example with 2 inputs (by default
# grad always takes partial derivative w.r.t. 1st input)

x = 1.5329423
y = -2.3

f = lambda x, y: jnp.exp(-jnp.pi*x**2) + 3*x**3 - jnp.cos(jnp.pi*y**2)
dfdx = grad(f) #same as grad(f, argnums=0), etc.
d2fdx2 = grad(dfdx) 
d3fdx3 = grad(d2fdx2) 
dfdy = grad(f, argnums=1)
d2fdy2 = grad(dfdy, argnums=1)

print(f(x, y), dfdx(x, y), d2fdx2(x, y), d3fdx3(x, y), dfdy(x, y), d2fdy2(x, y))

# Can also get df/dx and df/dy at the same time
grads = grad(f, argnums=(0, 1))
print(grads(x, y))

11.420368 21.143219 27.64676 17.557095 11.4187975 -132.96454
(Array(21.143219, dtype=float32, weak_type=True), Array(11.4187975, dtype=float32, weak_type=True))

In [ ]:

# JAX autodiff works not only for scalar fields, but also
# vector-valued functions (e.g. Jacobians)

from jax import jacfwd, jacrev
x = 1.
y = 1.
f = lambda x, y: x**2 + y**2 #paraboloid

#df/dx = 2x
#df/dy = 2y
#grad = [df/dx, df/dy]

#d2f/dx2 = 2
#d2f/dy2 = 2
#d2f/dxdy = d2df/dydx = 0
#H = [[d2f/dx2, d2fdxdy],[d2fdydx, d2fdy2]]

def hessian(f):
    return jit(jacfwd(jacrev(f, argnums=(0, 1)), argnums=(0, 1)))

print(jacrev(f, argnums=(0, 1))(x, y))
print(hessian(f)(x, y))

(Array(2., dtype=float32, weak_type=True), Array(2., dtype=float32, weak_type=True))
((Array(2., dtype=float32, weak_type=True), Array(0., dtype=float32, weak_type=True)), (Array(0., dtype=float32, weak_type=True), Array(2., dtype=float32, weak_type=True)))

In [ ]:

# Edge case of non-differentiable function

f = lambda x: abs(x)
visualize_fn(f)
dfdx = grad(f)
print(dfdx(0.)) # technically non-differentiable at x=0
# but JAX defaults to the derivative as x-->0+ from above

1.0

$\textbf{Problem}$: What does vmap() do in JAX?

$\textbf{Solution}$: Standing for vectorized map, it handles batch dimensions. The idea is one write functions as if they were to operate on a single data point, and vmap automatically transforms it to work on a batch of data points. This saves a lot of headache with either trying to write for loops or figure out complex broadcasting dimensions in NumPy.

In [ ]:

W = random.normal(key, (150, 100)) #e.g. treat as weights of linear NN layer 
batched_x = random.normal(key, (10, 100)) # e.g. a batch of 10 flattened 4x25  images
print(W)
print(batched_x)

def apply_matrix(x):
    return jnp.dot(W, x)

[[ 1.6226422   2.0252647  -0.43359444 ... -0.91352165  1.370097
  -0.7800775 ]
 [ 0.36481506  0.9761402  -0.0071727  ... -0.07060029  0.33603913
   2.354045  ]
 [-0.2431693   1.1728051   0.84588975 ... -0.7116935  -0.05395174
  -2.0926828 ]
 ...
 [ 0.6094885   0.8456557  -0.35652205 ... -1.2560906   0.42053804
   0.1999395 ]
 [-1.0437807   0.5227283   0.2781648  ...  0.83346254 -0.6578746
   0.67506284]
 [-0.3662531  -0.77806586  0.59587073 ...  1.1635658   0.475208
   0.6065461 ]]
[[ 1.62264216e+00  2.02526474e+00 -4.33594435e-01 -7.86173493e-02
   1.76090896e-01 -9.72089231e-01 -4.95298743e-01  4.94378597e-01
   6.64349318e-01 -9.50163484e-01  2.17953038e+00 -1.95515060e+00
   3.58570725e-01  1.57795131e-01  1.27708471e+00  1.51046479e+00
   9.70655978e-01  5.99608064e-01  2.47007050e-02 -1.91647720e+00
  -1.85934913e+00  1.72814405e+00  4.71903495e-02  8.14127982e-01
   1.31327674e-01  2.82847047e-01  1.24359429e+00  6.90280080e-01
  -8.00737441e-01 -7.40989983e-01 -1.53882873e+00  3.02691847e-01
  -2.07160451e-02  1.13287210e-01 -2.20654696e-01  7.05225617e-02
   8.53295803e-01 -8.21773827e-01 -1.46142114e-02 -1.50462165e-01
  -9.00135219e-01 -7.59072721e-01  3.33095133e-01  8.09249043e-01
   4.26925533e-02 -5.77671230e-01 -4.14398938e-01 -1.94125330e+00
   1.31611836e+00  7.54272819e-01  1.61709309e-01 -3.48330699e-02
  -1.33064091e+00  3.93620282e-01  4.82595831e-01  8.03829551e-01
  -6.33716822e-01  1.03875601e+00 -7.41591334e-01 -4.29958791e-01
  -2.25100428e-01 -5.19667149e-01 -1.66921651e+00  6.75354362e-01
   2.27387220e-01 -1.18004262e+00 -9.76733565e-01  1.19696045e+00
  -8.41275632e-01  6.59807801e-01  1.06801593e+00  3.15421283e-01
   4.37664032e-01  1.17185640e+00  9.07709897e-01  1.22262418e+00
  -5.46395242e-01  8.56304348e-01 -7.96577521e-03  4.73439127e-01
  -1.10903490e+00  2.64235139e+00  8.89576316e-01  9.95201647e-01
   2.55197197e-01  1.24961376e-01  1.16417289e+00  1.92963660e-01
  -1.90995440e-01 -4.36594725e-01 -1.14619887e+00  1.97602510e-01
   1.16866553e+00 -8.73398483e-01  8.81808579e-01 -3.44105691e-01
  -1.46149725e-01 -9.13521647e-01  1.37009704e+00 -7.80077517e-01]
 [ 3.64815056e-01  9.76140201e-01 -7.17270281e-03  2.10522056e-01
   1.90358415e-01  3.82912666e-01 -1.26563323e+00 -1.48435450e+00
  -1.14543624e-01  1.10371351e+00  1.98467016e-01  2.13889346e-01
  -6.60534799e-01 -7.27220058e-01  4.04439718e-01  1.89657375e-01
  -6.03179395e-01  9.45058823e-01  1.08387780e+00 -2.05607367e+00
  -7.13821530e-01  5.92868268e-01  1.05077624e+00 -1.46462381e+00
   6.60011351e-01 -3.01721781e-01  1.33131772e-01 -3.32813233e-01
   1.57000983e+00  5.74512124e-01  7.23415494e-01  6.96684480e-01
  -6.64234340e-01 -1.96695662e+00 -2.41625428e+00  2.73301542e-01
   1.16031730e+00  2.65512705e-01  6.90909326e-01 -2.56064266e-01
  -2.02274013e+00 -6.23128891e-01  2.79531687e-01 -1.35031724e+00
   1.01288453e-01  5.12681365e-01  2.64019489e-01 -1.82912755e+00
   1.43377745e+00  1.31885552e+00 -1.49532259e+00  9.33276117e-01
   1.40926480e+00 -1.67883754e-01 -1.18622862e-01 -2.42824897e-01
  -9.61759269e-01 -7.56359994e-01  2.57282567e+00 -1.06017923e+00
   3.12329054e-01  3.27511787e-01  8.28322321e-02 -1.08268857e+00
  -7.72234499e-01 -6.34604633e-01  1.22641027e+00 -1.48701501e+00
  -7.92869031e-01  5.53118527e-01 -1.18553972e+00  9.76909399e-01
  -4.38450336e-01 -3.29755992e-01  3.32547158e-01 -6.52719617e-01
  -1.20521224e+00 -8.86308253e-01 -2.10883713e+00 -1.55035362e-01
  -6.57932043e-01 -6.63254023e-01 -3.33620496e-02 -8.95929098e-01
   7.71168023e-02 -9.09823000e-01  1.27605200e+00 -4.01676625e-01
  -9.99925256e-01  1.73419788e-02  4.04541880e-01 -1.07132435e+00
   1.03666258e+00 -6.68480515e-01 -7.79318735e-02  1.20802212e+00
   2.00314546e+00 -7.06002936e-02  3.36039126e-01  2.35404491e+00]
 [-2.43169293e-01  1.17280507e+00  8.45889747e-01  7.28510857e-01
   5.97945750e-01 -7.66705036e-01  1.74257264e-01 -6.22359514e-01
   8.30672979e-01  3.31583172e-01 -8.89908016e-01  2.11538583e-01
  -9.24918503e-02 -1.61701906e+00 -4.12348479e-01  1.40828717e+00
   1.03313172e+00  1.08109426e+00 -1.14955103e+00 -3.21219079e-02
  -7.46680051e-02  5.98946154e-01 -1.19016111e+00 -6.29099309e-01
   5.55720888e-02  3.65632802e-01  5.46573997e-01  7.64500856e-01
  -3.39724004e-01  3.62569809e-01 -2.96992004e-01  1.03010583e+00
   9.65574086e-01  7.55760491e-01  1.21214367e-01  8.22795153e-01
   5.98026872e-01 -8.37458730e-01  5.43704927e-01 -3.22042048e-01
  -4.44730967e-01  9.14655566e-01  1.21259427e+00  1.79783368e+00
  -1.60528362e+00  7.97706470e-02 -7.69092619e-01 -1.94633031e+00
   1.59343791e+00 -9.19291854e-01 -1.71621680e+00  1.11879945e+00
   2.03491643e-01 -1.80223882e-01  1.74956903e-01 -7.34732687e-01
   4.10009056e-01 -1.53308785e+00 -1.70366681e+00  2.04102325e+00
  -2.04918474e-01 -8.45943689e-02  6.67714596e-01 -1.38450539e+00
   3.79239351e-01 -7.40188420e-01  1.30270433e+00 -1.02559209e+00
   4.96420115e-01 -7.00343609e-01 -3.66578884e-02 -4.15897846e-01
  -2.48955950e-01  9.09369051e-01 -1.88014939e-01 -7.17376173e-01
   3.89687389e-01 -2.08640948e-01 -7.79499948e-01  1.29690838e+00
  -9.13590670e-01  1.33042842e-01  8.17202270e-01 -4.76543754e-01
   4.17642653e-01 -9.76111740e-03 -9.28037405e-01  6.50380313e-01
  -3.19879699e+00  4.65263397e-01  3.35585117e-01  2.14436688e-02
  -2.76631832e-01  4.24526855e-02 -3.00587714e-01 -6.31640971e-01
  -1.22646141e+00 -7.11693525e-01 -5.39517365e-02 -2.09268284e+00]
 [ 4.73754674e-01 -1.59625351e+00 -1.53968468e-01 -2.44954348e+00
   6.49740636e-01 -7.07279503e-01  1.60122812e-01 -1.56832671e+00
   2.35200658e-01  1.16291511e+00  4.25386399e-01 -3.92973781e-01
   7.57110953e-01 -1.14584528e-01 -1.07421279e+00 -6.47926211e-01
  -9.07325804e-01  2.22105756e-01  3.95380974e-01  1.78448841e-01
  -5.79642057e-02  1.04901743e+00  1.34441698e+00 -9.69599932e-02
   9.23476666e-02  8.63941193e-01  1.11614284e-03 -1.05547857e+00
   1.09360147e+00 -1.33676386e+00  8.93752217e-01  3.83038670e-02
  -2.31079721e+00  2.61638403e-01  5.76163709e-01  7.12172329e-01
   1.13079399e-01 -1.81492853e+00  1.04789495e+00 -1.23274231e+00
   5.07376343e-02 -1.33968818e+00  2.26974034e+00  2.03760728e-01
   5.24858892e-01 -1.25730753e+00 -4.17512745e-01  6.30214334e-01
  -5.04626036e-01  2.55590463e+00  5.74609153e-02  4.38864678e-01
  -2.43704557e-01  1.28732228e+00  1.14161777e+00  5.20770967e-01
   2.16976404e-01  1.14142621e+00  8.94764543e-01  1.80909836e+00
  -1.90802503e+00 -9.99371231e-01  7.12768495e-01 -6.34517670e-01
  -1.24687783e-01  6.85439765e-01  9.78687882e-01 -2.72447288e-01
  -1.24981749e+00 -8.90857697e-01  1.01692431e-01  9.47433591e-01
   9.78131071e-02  1.25865614e+00  1.04573154e+00 -1.85578093e-01
   5.13851285e-01 -1.96585572e+00 -6.18929446e-01 -1.59196496e+00
  -7.94556320e-01  1.35239506e+00  1.30747497e-01  1.08340538e+00
   7.94674933e-01 -1.70419037e-01 -1.07815914e-01 -8.67752433e-01
   2.74028946e-02  1.42733014e+00 -3.35451245e-01 -1.18081085e-02
  -7.14428008e-01  1.29715097e+00 -1.53426766e+00 -7.59699047e-01
   1.63791358e-01 -1.65026736e+00  1.48562670e+00 -1.72110963e+00]
 [-3.48294787e-02  2.26507974e+00  4.86799538e-01  8.02250922e-01
  -5.03105819e-01 -1.51174557e+00 -1.23069298e+00 -4.37387377e-01
  -1.19922660e-01 -1.32741228e-01  9.81089592e-01 -2.00786293e-01
   8.46012473e-01  7.58226991e-01  1.37318218e+00  4.99171913e-01
   7.41392896e-02  1.14247692e+00  1.19223285e+00  1.54212606e+00
  -2.56369328e+00 -5.90210617e-01 -1.69458699e+00  3.01133752e-01
   9.47523534e-01 -1.28467157e-01 -6.07066035e-01 -6.00612879e-01
  -3.29539704e+00  3.31878364e-01  3.49914461e-01 -8.02474558e-01
   9.00005996e-02  8.55811656e-01 -6.35393858e-01 -1.88428596e-01
   8.55382442e-01 -7.71138191e-01  1.48878348e+00 -5.60926378e-01
   6.07925117e-01  1.13047779e+00 -1.20606339e+00  8.76124084e-01
   3.93461764e-01  5.42526424e-01  9.87375915e-01  1.54771483e+00
  -1.11704183e+00  9.66534257e-01 -1.50813270e+00 -4.18101907e-01
  -8.16014290e-01 -6.96839750e-01 -2.56920576e-01  4.37291175e-01
  -2.53834738e-03  2.33468676e+00 -1.50986416e-02 -6.59652352e-01
  -4.65559602e-01 -1.34911144e+00  4.10299897e-02  4.49487656e-01
  -3.28320116e-01  5.82158938e-03 -9.17235196e-01  5.47479868e-01
   1.00075707e-01  5.53778887e-01  7.95560360e-01  5.82074165e-01
   4.40597296e-01  1.41883627e-01 -1.60682452e+00  8.55871812e-02
   4.72201049e-01  6.77530110e-01 -1.82441995e-01  1.01498105e-01
  -8.04406777e-03  1.96414077e+00 -1.11661828e+00 -1.56128228e+00
  -4.74351019e-01  1.91607726e+00 -1.19328606e+00 -2.86463916e-01
   1.25017092e-01  1.13316298e+00 -2.74157941e-01 -1.64727375e-01
  -1.72950840e+00  1.12391174e+00 -6.66688561e-01  7.28739873e-02
   1.04647315e+00 -1.55617642e+00 -2.54764885e-01 -8.18689287e-01]
 [ 8.91156673e-01  3.18052262e-01  2.09722710e+00  1.73778522e+00
  -1.41319370e+00 -7.47485280e-01  6.66629076e-01  7.39560783e-01
   8.95429015e-01  1.51626244e-01  1.38383520e+00 -1.36733592e-01
   8.11541319e-01 -4.33953226e-01  7.94477880e-01 -6.55612409e-01
  -1.38271320e+00 -3.39266628e-01 -2.93511152e-01  1.46516097e+00
   1.71347570e+00 -1.40051532e+00  1.12666571e+00 -4.97829884e-01
  -5.93189120e-01 -1.45751333e+00 -2.71108794e+00 -2.00019574e+00
  -9.48191524e-01 -4.44346368e-01  8.13890696e-01 -7.42613897e-02
   9.32008401e-02 -1.23251045e+00 -8.50963406e-03  2.05505773e-01
  -1.48191705e-01 -9.23814103e-02 -2.20937908e-01  2.62501180e-01
   1.08574092e-01  7.75018990e-01  1.11870992e+00  7.86271334e-01
   9.40064609e-01  1.32092059e+00 -5.05305588e-01  1.13251865e+00
   1.40149787e-01 -1.42711639e+00 -2.81242430e-01 -1.66696763e+00
   1.10434997e+00  1.75778019e+00 -1.71847045e+00 -2.79099345e-01
   1.51310873e+00 -8.37235510e-01 -7.63156950e-01 -1.06069550e-01
   2.34096125e-01 -2.71091008e+00  2.44496793e-01  1.89915252e+00
   8.01037431e-01  1.23803353e+00 -9.71419290e-02 -1.92206979e+00
   8.28902960e-01 -4.83347476e-01  5.25316775e-01  2.30684206e-02
  -5.12029767e-01 -7.48020053e-01  4.11946565e-01  5.49255192e-01
  -2.23095298e-01 -5.12054920e-01  2.52339333e-01 -9.49054807e-02
   1.07594228e+00 -9.24595118e-01 -6.69982493e-01 -1.06690586e+00
   6.30436301e-01  2.94996917e-01  1.32940316e+00 -1.39383161e+00
  -5.40680647e-01 -7.85788894e-01  3.04403752e-01  4.66728628e-01
   2.67849475e-01  1.44306827e+00 -2.27237716e-01  1.02745861e-01
  -5.52183032e-01 -5.09359360e-01  1.01352119e+00  6.09791458e-01]
 [ 3.99991989e-01 -6.25152647e-01  7.01872110e-01 -1.16544163e+00
   8.80496144e-01  7.28263557e-01  3.84819478e-01 -1.36815324e-01
  -4.19357091e-01 -2.99014002e-01  6.20382190e-01 -1.30659953e-01
   1.17253888e+00 -5.22038400e-01  1.64784598e+00  8.08855772e-01
   1.11875927e+00  1.02307034e+00  1.02533817e+00 -8.38138282e-01
  -6.32627964e-01  8.43035996e-01 -4.72003847e-01 -1.91650724e+00
   1.87835085e+00  6.62476659e-01 -1.64310064e-03  1.54904974e+00
  -4.31367248e-01 -9.94612157e-01 -2.51327772e-02 -8.24098051e-01
   5.06091237e-01 -6.20458759e-02 -3.90430093e-01 -9.47100639e-01
  -1.86841384e-01 -8.28972280e-01 -1.91550240e-01 -2.12299919e+00
   8.34999502e-01 -2.38999426e-01 -2.41668820e-01 -9.00301933e-01
  -2.25573361e-01  1.54047862e-01  9.26408947e-01 -2.11789295e-01
   4.12847757e-01  1.14952362e+00 -7.36300349e-01  2.04192728e-01
   1.28266513e+00 -6.70982182e-01 -1.83460546e+00  2.15686157e-01
   7.07491457e-01 -5.47649086e-01  1.27539206e+00  5.64734936e-01
   2.26897240e-01 -1.17738798e-01 -4.27806437e-01  2.72378884e-02
   1.63076770e+00 -4.90351528e-01  5.36693037e-01  8.68616879e-01
   8.11795175e-01  3.30178648e-01  2.17748213e+00 -7.27532327e-01
  -1.45456925e-01 -1.59931540e+00 -2.77693480e-01 -4.46349621e-01
  -9.30218279e-01 -2.19141650e+00  8.11784804e-01  7.18652382e-02
   6.57420337e-01  1.03664204e-01 -1.04592490e+00  6.97659552e-01
   1.92154855e-01  6.56601965e-01  8.16967726e-01 -1.21444792e-01
   7.82743096e-01  1.41578543e+00  3.41761321e-01  1.23176627e-01
   6.58716798e-01 -9.68537092e-01  8.28790963e-01 -6.14581816e-02
   4.16769356e-01  7.93242037e-01  1.27035511e+00 -2.98217845e+00]
 [ 1.41787744e+00  4.30863053e-01 -1.35842001e+00 -2.31687352e-01
   4.81799468e-02 -1.13010384e-01 -2.81411290e-01  1.23678505e+00
  -4.91829105e-02  3.90700817e-01 -1.18719149e+00  5.87880611e-01
   1.05693567e+00  1.75502792e-01 -8.20879996e-01 -6.93775490e-02
   1.40034509e+00 -3.20545942e-01 -8.02284122e-01 -1.90450406e+00
   1.86188713e-01 -6.00135505e-01 -3.11949283e-01 -1.58620250e+00
   4.76210535e-01 -1.73004854e+00 -6.93736970e-01 -5.08554578e-01
   3.16312289e+00  1.95174009e-01 -2.76474714e-01  1.33630812e+00
   2.34854460e-01 -2.47333512e-01  7.94547141e-01  8.25466216e-01
   5.37205935e-01  9.24634635e-01 -1.38021624e+00 -2.07926393e-01
  -4.06866372e-01 -1.34310174e+00  1.07187748e+00  4.51775134e-01
  -9.03886437e-01  7.52393663e-01 -5.68250000e-01  6.33994997e-01
   1.08135533e+00 -1.00676358e+00 -9.19247627e-01 -2.00193954e+00
  -1.03497982e+00 -3.54129493e-01 -1.17053652e+00  2.89141273e+00
   2.00288579e-01  2.43504629e-01  5.16683877e-01 -1.94491204e-02
  -1.23140061e+00  5.12821823e-02 -4.50057536e-01 -1.03455268e-01
   2.95981437e-01  9.00486037e-02 -1.65372467e+00 -7.55375445e-01
   4.64709401e-01 -6.81560993e-01 -1.21171749e+00 -1.10982382e+00
  -8.82262439e-02 -4.89559084e-01  3.21174353e-01  4.49467182e-01
  -1.44262266e+00 -6.77745342e-02 -4.52509262e-02 -4.69268233e-01
   1.16332404e-01  1.82556927e+00  6.26976192e-01 -9.81169581e-01
   4.06791836e-01 -1.66363668e+00  7.90482312e-02 -1.15621150e+00
   1.21551013e+00 -1.50499329e-01 -8.60824715e-03  1.26619697e+00
  -6.30170465e-01 -9.19624329e-01  4.25071388e-01  1.68217826e+00
  -1.06247319e-02 -2.80735463e-01 -3.46607894e-01  6.16329730e-01]
 [ 1.23951900e+00 -1.29802644e+00 -1.11810875e+00 -1.61887482e-01
   1.66484877e-01  4.02306058e-02  2.34848931e-01 -6.17022276e-01
  -1.01548862e+00 -1.59668422e+00 -4.61611271e-01  1.00245929e+00
   2.30857298e-01  2.06788397e+00 -6.77438915e-01 -1.82206213e-01
   2.47246489e-01 -1.72794685e-01 -1.61421788e+00  3.54998648e-01
   1.22171605e+00  1.04115450e+00  6.81458652e-01 -6.77761137e-01
   1.90401602e+00  1.37473094e+00  1.93106627e+00 -8.90373051e-01
   2.69743651e-01  1.10350811e+00  9.85456467e-01 -4.17405128e-01
   1.45321739e+00 -1.62904516e-01  2.61666441e+00 -1.13660252e+00
  -5.23015976e-01  1.41417384e+00 -1.45638537e+00  1.39924064e-01
  -1.81835264e-01 -1.85092092e+00  6.44125581e-01 -3.70008737e-01
   3.54113102e-01  1.47451103e+00 -1.08801985e+00  2.43944931e+00
  -1.32531929e+00 -1.80648351e+00 -4.23250198e-01  7.63600230e-01
  -5.85210204e-01  2.35295343e+00  1.04250467e+00  1.02737755e-01
  -2.87180662e-01 -8.28941539e-02  6.05848372e-01  2.33335465e-01
   1.96693033e-01 -7.93169662e-02  1.89767814e+00 -2.49023829e-03
   9.39399183e-01  1.27057350e+00 -1.26574382e-01  1.88864917e-01
   2.41839439e-01 -5.35686433e-01 -1.35840702e+00 -1.48494470e+00
   1.21456063e+00  3.16799998e-01  7.81244099e-01  1.40346968e+00
  -8.55279490e-02 -4.50297594e-01 -1.24199247e+00 -5.65499604e-01
   1.21995819e+00 -5.08152306e-01  6.65360019e-02  1.64739639e-01
   1.54252136e+00  1.00633180e+00 -1.25837934e+00  5.13076067e-01
   4.77847643e-02  4.05105203e-02 -6.82877123e-01 -1.07491016e-01
   5.98556817e-01 -2.11435556e+00  1.08828628e+00  1.69309035e-01
  -1.96819067e+00  1.22727036e+00 -1.31954658e+00 -3.26137334e-01]
 [ 7.08347976e-01  8.58157128e-02  8.75563264e-01  2.74913341e-01
   4.60563987e-01  3.63558292e-01 -3.28417599e-01  1.39410818e+00
   2.12482005e-01  3.25346977e-01 -5.32687724e-01  8.93540680e-02
   5.19629002e-01 -1.36936378e+00 -6.61399513e-02  1.50276494e+00
   1.28715611e+00 -2.44274974e-01  5.93465090e-01  5.05267680e-01
   6.96502507e-01 -1.56812537e+00  1.53453970e+00  8.01590204e-01
  -1.70177698e+00 -5.72526097e-01 -7.42326021e-01 -4.66878384e-01
  -2.30425686e-01 -3.67261142e-01  1.15874124e+00 -2.89703548e-01
  -2.18252748e-01  3.10105979e-01 -1.29287863e+00  2.05113482e+00
   2.57617265e-01 -1.06610370e+00 -5.81136525e-01  1.32945165e-01
   6.66835070e-01  7.49220848e-01  5.12089491e-01  4.30191994e-01
  -1.51903883e-01  7.75620282e-01  1.16832733e+00  4.68485653e-01
  -4.53565627e-01 -9.55212057e-01 -2.65528321e-01 -3.27313989e-01
   9.56872046e-01  5.94931424e-01 -3.91686225e+00  8.41580987e-01
  -4.13113117e-01 -1.14385709e-01 -1.87210631e+00  5.46258092e-02
  -5.05993187e-01 -1.67221451e+00  2.15138698e+00 -2.31907427e-01
  -3.46957266e-01 -3.51842254e-01  4.06620145e-01  7.34579444e-01
   1.59595609e+00 -1.41492343e+00 -9.70596373e-01  8.89268756e-01
  -7.70626605e-01 -1.15664482e+00 -1.54152620e+00  2.41009980e-01
  -2.38504753e-01 -2.03631926e+00 -1.29285181e+00  3.12330008e-01
  -2.23036790e+00 -1.19731987e+00 -1.44340980e+00 -1.25551546e+00
  -1.41780686e+00 -6.98669374e-01 -1.49969471e+00 -4.50949132e-01
  -1.64870260e-04  8.19425404e-01 -6.81385696e-01 -7.35609353e-01
  -6.61660969e-01 -1.59738159e+00 -2.30075568e-01  1.22990265e-01
   1.80285692e-01 -5.41782714e-02 -1.50115645e+00 -2.67956465e-01]]

In [ ]:

def naively_batched_apply_matrix(batched_x):
    return jnp.stack([apply_matrix(x) for x in batched_x])

print("Naively batched time:")
%timeit naively_batched_apply_matrix(batched_x).block_until_ready( )

Naively batched time:
1.77 ms ± 93.3 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [ ]:

@jit
def batched_apply_matrix(batched_x):
    return jnp.dot(batched_x, W.T) #had to completely rewrite function to vectorize it

print("Manually batched time:")
%timeit batched_apply_matrix(batched_x).block_until_ready()

Manually batched time:
223 μs ± 49.3 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [ ]:

@jit # note: one can arbitrarily compose JAX transformations
def vmap_batched_apply_matrix(batched_x):
    return vmap(apply_matrix)(batched_x) #much simpler! and very efficient!

print("Automatically vectorized with vmap time:")
%timeit vmap_batched_apply_matrix(batched_x).block_until_ready()

Automatically vectorized with vmap time:
157 μs ± 18.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [ ]:

def apply_matrix(W, x): #accept both W and x args, in line with func prog paradigm
    return jnp.dot(W, x) 

@jit # note: we can arbitrarily compose JAX transformations
def vmap_batched_apply_matrix(W, batched_x):
    return vmap(apply_matrix)(W, batched_x) #much simpler! and very efficient!

print("Automatically vectorized with vmap time:")
%timeit vmap_batched_apply_matrix(W, batched_x).block_until_ready()
# crashes because W does not have a batch dimension.

Automatically vectorized with vmap time:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[167], line 9
      6     return vmap(apply_matrix)(W, batched_x) #much simpler! and very efficient!
      8 print("Automatically vectorized with vmap time:")
----> 9 get_ipython().run_line_magic('timeit', 'vmap_batched_apply_matrix(W, batched_x).block_until_ready()')

File ~/jax_linus/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2482, in InteractiveShell.run_line_magic(self, magic_name, line, _stack_depth)
   2480     kwargs['local_ns'] = self.get_local_scope(stack_depth)
   2481 with self.builtin_trap:
-> 2482     result = fn(*args, **kwargs)
   2484 # The code below prevents the output from being displayed
   2485 # when using magics with decorator @output_can_be_silenced
   2486 # when the last Python token in the expression is a ';'.
   2487 if getattr(fn, magic.MAGIC_OUTPUT_CAN_BE_SILENCED, False):

File ~/jax_linus/lib/python3.10/site-packages/IPython/core/magics/execution.py:1209, in ExecutionMagics.timeit(self, line, cell, local_ns)
   1207 for index in range(0, 10):
   1208     number = 10 ** index
-> 1209     time_number = timer.timeit(number)
   1210     if time_number >= 0.2:
   1211         break

File ~/jax_linus/lib/python3.10/site-packages/IPython/core/magics/execution.py:174, in Timer.timeit(self, number)
    172 gc.disable()
    173 try:
--> 174     timing = self.inner(it, self.timer)
    175 finally:
    176     if gcold:

File <magic-timeit>:1, in inner(_it, _timer)

    [... skipping hidden 14 frame]

Cell In[167], line 6, in vmap_batched_apply_matrix(W, batched_x)
      4 @jit # note: we can arbitrarily compose JAX transformations
      5 def vmap_batched_apply_matrix(W, batched_x):
----> 6     return vmap(apply_matrix)(W, batched_x)

    [... skipping hidden 2 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/api.py:1248, in _mapped_axis_size(fn, tree, vals, dims, name)
   1246   else:
   1247     msg.append(f"  * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n")
-> 1248 raise ValueError(''.join(msg)[:-2])

ValueError: vmap got inconsistent sizes for array axes to be mapped:
  * one axis had size 150: axis 0 of argument W of type float32[150,100];
  * one axis had size 10: axis 0 of argument x of type float32[10,100]

In [ ]:

# Solution: use the in_axes arg for vmap

def apply_matrix(W, x): #accept both W and x args, in line with func prog paradigm
    return jnp.dot(W, x) 

@jit # note: we can arbitrarily compose JAX transformations
def vmap_batched_apply_matrix(W, batched_x):
    return vmap(apply_matrix, in_axes=(None, 0))(W, batched_x) #much simpler! and very efficient!

print("Automatically vectorized with vmap time:")
%timeit vmap_batched_apply_matrix(W, batched_x).block_until_ready()

Automatically vectorized with vmap time:
197 μs ± 29.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

$\textbf{Problem}$: Why is JAX’s API layering similar to an onion?

$\textbf{Solution}$: The highest layer of the API is NumPy, followed by Lax, followed by the XLA compiler. The Lax API is stricter than NumPy but also more powerful; it’s a Python wrapper around XLA.

In [ ]:

# Example 1: Lax is stricter than NumPy

print(jnp.add(1, 1.0))
print(lax.add(1, 1.0))

2.0

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[171], line 4
      1 # Example 1: Lax is stricter than NumPy
      3 print(jnp.add(1, 1.0))
----> 4 print(lax.add(1, 1.0))

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/lax/lax.py:1196, in add(x, y)
   1176 r"""Elementwise addition: :math:`x + y`.
   1177 
   1178 This function lowers directly to the `stablehlo.add`_ operation.
   (...)
   1193 .. _stablehlo.add: https://openxla.org/stablehlo/spec#add
   1194 """
   1195 x, y = core.standard_insert_pvary(x, y)
-> 1196 return add_p.bind(x, y)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:536, in Primitive.bind(self, *args, **params)
    534 def bind(self, *args, **params):
    535   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> 536   return self._true_bind(*args, **params)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:552, in Primitive._true_bind(self, *args, **params)
    550 trace_ctx.set_trace(eval_trace)
    551 try:
--> 552   return self.bind_with_trace(prev_trace, args, params)
    553 finally:
    554   trace_ctx.set_trace(prev_trace)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:562, in Primitive.bind_with_trace(self, trace, args, params)
    559   with set_current_trace(trace):
    560     return self.to_lojax(*args, **params)  # type: ignore
--> 562 return trace.process_primitive(self, args, params)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:1066, in EvalTrace.process_primitive(self, primitive, args, params)
   1064 args = map(full_lower, args)
   1065 check_eval_args(args)
-> 1066 return primitive.impl(*args, **params)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/dispatch.py:91, in apply_primitive(prim, *args, **params)
     89 prev = lib.jax_jit.swap_thread_local_state_disable_jit(False)
     90 try:
---> 91   outs = fun(*args)
     92 finally:
     93   lib.jax_jit.swap_thread_local_state_disable_jit(prev)

    [... skipping hidden 24 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/lax/lax.py:8799, in check_same_dtypes(name, *avals)
   8797   equiv = _JNP_FUNCTION_EQUIVALENTS[name]
   8798   msg += f" (Tip: jnp.{equiv} is a similar function that does automatic type promotion on inputs)."
-> 8799 raise TypeError(msg.format(name, ", ".join(str(a.dtype) for a in avals)))

TypeError: lax.add requires arguments to have the same dtypes, got int32, float32. (Tip: jnp.add is a similar function that does automatic type promotion on inputs).

In [ ]:

# Example 2: Lax is more powerful (tradeoff: less user-friendly)

x = jnp.array([1, 2, 1])
y = jnp.ones(10)

# NumPy API
convolution_jnp = jnp.convolve(x, y)

# Lax API
convolution_lax = lax.conv_general_dilated(
    x.reshape(1, 1, 3).astype(float),
    y.reshape(1, 1, 10),
    window_strides=(1,),
    padding=[(len(y)-1, len(y)-1)]
)

print(convolution_jnp)
print(convolution_lax)
print(convolution_lax[0][0]) # returns batched result, hence need for this indexing
# see:

[1. 3. 4. 4. 4. 4. 4. 4. 4. 4. 3. 1.]
[[[1. 3. 4. 4. 4. 4. 4. 4. 4. 4. 3. 1.]]]
[1. 3. 4. 4. 4. 4. 4. 4. 4. 4. 3. 1.]

In [ ]:

# another JIT example
def norm(X):
    X = X - X.mean(0)
    return X/X.std(0)

norm_compiled = jit(norm)

X = random.normal(key, (10000, 100), dtype=jnp.float32)
%timeit norm(X).block_until_ready()
%timeit norm_compiled(X).block_until_ready()

1.9 ms ± 153 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.91 ms ± 52.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [ ]:

# Example of a failure: array shapes must be static.
def get_negatives(x):
    return x[x < 0]

x = random.normal(key, (10,), dtype=jnp.float32)
print(get_negatives(x))

[-0.43359444 -0.07861735 -0.97208923 -0.49529874 -0.9501635 ]

In [ ]:

# but this fails:
print(jit(get_negatives)(x))

---------------------------------------------------------------------------
NonConcreteBooleanIndexError              Traceback (most recent call last)
Cell In[199], line 2
      1 # but this fails:
----> 2 print(jit(get_negatives)(x))

    [... skipping hidden 14 frame]

Cell In[198], line 3, in get_negatives(x)
      2 def get_negatives(x):
----> 3     return x[x < 0]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:1083, in _forward_operator_to_aval.<locals>.op(self, *args)
   1082 def op(self, *args):
-> 1083   return getattr(self.aval, f"_{name}")(self, *args)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:657, in _getitem(self, item)
    656 def _getitem(self, item):
--> 657   return indexing.rewriting_take(self, item)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:645, in rewriting_take(arr, idx, indices_are_sorted, unique_indices, mode, fill_value, out_sharding)
    639     if (isinstance(aval, core.DShapedArray) and aval.shape == () and
    640         dtypes.issubdtype(aval.dtype, np.integer) and
    641         not dtypes.issubdtype(aval.dtype, dtypes.bool_) and
    642         isinstance(arr.shape[0], int)):
    643       return lax.dynamic_index_in_dim(arr, idx, keepdims=False)
--> 645 treedef, static_idx, dynamic_idx = split_index_for_jit(idx, arr.shape)
    646 internal_gather = partial(
    647     _gather, treedef=treedef, static_idx=static_idx,
    648     indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
    649     mode=mode, fill_value=fill_value)
    650 if out_sharding is not None:

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:738, in split_index_for_jit(idx, shape)
    734   raise TypeError(f"JAX does not support string indexing; got {idx=}")
    736 # Expand any (concrete) boolean indices. We can then use advanced integer
    737 # indexing logic to handle them.
--> 738 idx = _expand_bool_indices(idx, shape)
    740 leaves, treedef = tree_flatten(idx)
    741 dynamic = [None] * len(leaves)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/indexing.py:1059, in _expand_bool_indices(idx, shape)
   1055   abstract_i = core.get_aval(i)
   1057 if not core.is_concrete(i):
   1058   # TODO(mattjj): improve this error by tracking _why_ the indices are not concrete
-> 1059   raise errors.NonConcreteBooleanIndexError(abstract_i)
   1060 elif np.ndim(i) == 0:
   1061   out.append(bool(i))

NonConcreteBooleanIndexError: Array boolean indices must be concrete; got bool[10]

See https://docs.jax.dev/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError

In [ ]:

# So why does this happen? --> tracing on different levels of abstraction

@jit 
def f(x, y):
    print("Running f().")
    print(f"x={x}")
    print(f"y={y}")
    result = jnp.dot(x+1, y+1)
    print(f"result = {result}")
    return result

x = np.random.randn(3, 4)
y = np.random.randn(4)
print(f(x, y))

x2 = np.random.randn(3, 4)
y2 = np.random.randn(4)
print("Second Call (notice no more print statement side effects):")
print(f(x2, y2))

# Note: any time one puts the same shapes + types into a jitted function, fast b/c tracer object already XLA-compiled
# S + T! shapes + types!
x3 = np.random.randn(3, 5)
y3 = np.random.randn(5)
print(f(x3, y3)) # notice how jit has to retrace it

Running f().
x=Traced<float32[3,4]>with<DynamicJaxprTrace>
y=Traced<float32[4]>with<DynamicJaxprTrace>
result = Traced<float32[3]>with<DynamicJaxprTrace>
[1.0699788 4.141973  2.184647 ]
Second Call (notice no more print statement side effects):
[9.221527  5.545151  2.8010619]
Running f().
x=Traced<float32[3,5]>with<DynamicJaxprTrace>
y=Traced<float32[5]>with<DynamicJaxprTrace>
result = Traced<float32[3]>with<DynamicJaxprTrace>
[5.1654205 4.1818676 6.113046 ]

In [ ]:

def f(x, y):
    # same function as above but w/o the print statements, as would be written in practice
    return jnp.dot(x + 1, y + 1)

print(make_jaxpr(f)(x, y))
# produces a flow model showing what jit creates in the background when tracing
# can go to JAX docs to get better understanding of this grammar

{ lambda ; a:f32[3,4] b:f32[4]. let
    c:f32[3,4] = add a 1.0:f32[]
    d:f32[4] = add b 1.0:f32[]
    e:f32[3] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=float32
    ] c d
  in (e,) }

In [ ]:

@jit
def f(x, neg):
    return -x if neg else X

f(1, True)

---------------------------------------------------------------------------
TracerBoolConversionError                 Traceback (most recent call last)
Cell In[14], line 5
      1 @jit
      2 def f(x, neg):
      3     return -x if neg else X
----> 5 f(1, True)

    [... skipping hidden 14 frame]

Cell In[14], line 3, in f(x, neg)
      1 @jit
      2 def f(x, neg):
----> 3     return -x if neg else X

    [... skipping hidden 1 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:1661, in concretization_function_error.<locals>.error(self, arg)
   1660 def error(self, arg):
-> 1661   raise TracerBoolConversionError(arg)

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function f at /tmp/ipykernel_424021/236938057.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument neg.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError

$\textbf{Problem}$: Why does running the above code cell lead to the error message shown? What can be done to debug this error?

$\textbf{Solution}$: When you jit a function, there is always a $2$-step procedure:

Tracing (Recipe Phase): JAX runs the function once using placeholder objects called tracers instead of the actual data. This is meant purely to record the computational graph of the function.
Execution (Cooking Phase): JAX compiles the graph into XLA code to run on GPU/TPU.

So the conflict here is that standard Python control flow (such as the if/else conditional statement in the function above) happened during the $1^{\text{st}}$ step of the jit compilation. The Python interpreter would have checked whether neg is True or False, but at the time it was still a tracer object, which basically told the interpreter “I don’t know what value I hold yet; I represent a value that will exist later when we run the XLA code”.

To debug this error, one can instead make neg a static argument (see below).

In [ ]:

from functools import partial

@partial(jit, static_argnums=(1, ))
def f(x, neg):
    print(x)
    return -x if neg else x

print(f(1, True))
print(f(2, True)) # no more tracing here b/c still True which was already traced
print(f(2, False))
print(f(23, False)) # same comment here

Traced<~int32[]>with<DynamicJaxprTrace>
-1
-2
Traced<~int32[]>with<DynamicJaxprTrace>
2
23

In [39]:

@jit
def f(x):
    print(x) #tracer object
    print(x.shape) #concrete value
    print(jnp.array(x.shape).prod()) #tracer object again
    return x.reshape(jnp.array(x.shape).prod()) #fail b/c try to apply reshape to a tracer

x = jnp.ones((2, 3))
print(f(x))

Traced<float32[2,3]>with<DynamicJaxprTrace>
(2, 3)
Traced<int32[]>with<DynamicJaxprTrace>

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[39], line 9
      6     return x.reshape(jnp.array(x.shape).prod()) #fail b/c try to apply reshape to a tracer
      8 x = jnp.ones((2, 3))
----> 9 print(f(x)) 

    [... skipping hidden 14 frame]

Cell In[39], line 6, in f(x)
      4 print(x.shape) #concrete value
      5 print(jnp.array(x.shape).prod()) #tracer object again
----> 6 return x.reshape(jnp.array(x.shape).prod())

    [... skipping hidden 2 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py:454, in _compute_newshape(arr, newshape)
    452 except:
    453   newshape = [newshape]
--> 454 newshape = core.canonicalize_shape(newshape)  # type: ignore[arg-type]
    455 neg1s = [i for i, d in enumerate(newshape) if type(d) is int and d == -1]
    456 if len(neg1s) > 1:

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/core.py:1864, in canonicalize_shape(shape, context)
   1862 except TypeError:
   1863   pass
-> 1864 raise _invalid_shape_error(shape, context)

TypeError: Shapes must be 1D sequences of concrete values of integer type, got [Traced<int32[]>with<DynamicJaxprTrace>].
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
The error occurred while tracing the function f at /tmp/ipykernel_424021/3415433522.py:1 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:i32[] = reduce_prod[axes=(0,)] b
    from line /tmp/ipykernel_424021/3415433522.py:6 (f)

In [ ]:

@jit
def f(x):
    return x.reshape(np.prod(x.shape)) #np.prod works b/c it returns concrete array rather than tracer object

print(f(x))

[1. 1. 1. 1. 1. 1.]

$\textbf{Problem}$: What does it mean that JAX is only designed to work with “pure functions”?

$\textbf{Solution}$: Informally, a function is pure iff:

All inputs are passed through the function arguments, all outputs are passed through the return statement of the function.
When passed with the same input arguments, the same output result is always obtained.

In [46]:

# Example: identity function

def impure_print_side_effect(x):
    print("Executing function:") #violates #1
    return x

print("First call (tracing)", jit(impure_print_side_effect)(4))
print("Second call, no more side effects b/c traced and XLA-compiled already", jit(impure_print_side_effect)(5))
print("Third call, need to trace computation graph again, type change:", jit(impure_print_side_effect)(jnp.array([5,])))

Executing function:
First call (tracing) 4
Second call, no more side effects b/c traced and XLA-compiled already 5
Executing function:
Third call, need to trace computation graph again, type change: [5]

In [52]:

# Example: don't intefere with global variable values!

g = 0
def impure_uses_globals(x):
    return x + g

print("First call (caches value of g=0):", jit(impure_uses_globals)(4))

# Now do the crime of updating a global variable:
g = 10

print("Second call looks wrong:", jit(impure_uses_globals)(5))

# If change the type of the func arg, JAX will retrace and now it will read the
# most recent value of the global variable

print("Third call:", jit(impure_uses_globals)(jnp.array([4])))

First call (caches value of g=0): 4
Second call looks wrong: 5
Third call: [14]

In [59]:

# Example: valid pure function (Haiku/Flax built on this idea)

def pure_uses_internal_state(x):
    state = dict(even=0, odd=0)
    for i in range(10):
        state["even" if i%2==0 else "odd"] += x
    return state["even"], state["odd"]

print(jit(pure_uses_internal_state)(7))

(Array(35, dtype=int32, weak_type=True), Array(35, dtype=int32, weak_type=True))

In [ ]:

# Example: iterators are forbidden!

# lax.fori_loop: similarly for lax.scan, lax.cond, etc.
array = jnp.arange(10)
print("Correct Answer:", lax.fori_loop(0, 10, lambda i, x: x+array[i], 0))

# however this one breaks b/c an iterator is used.
# Iterators are stateful objects, so they violate the purity constraint of JAX functions
iterator = iter(range(10))
print("Incorrect Answer:", lax.fori_loop(0, 10, lambda i,x: x+next(iterator), 0))

Correct Answer: 45
Incorrect Answer: 0

$\text{Problem}$: Given a JAX NumPy array, how should one update the value of an element in the array in-place?

$\text{Solution}$: Using the .at[].set() syntax, for example:

In [65]:

jax_arr = jnp.zeros(shape=(3, 3), dtype=jnp.float32)
updated_jax_arr = jax_arr.at[1,:].set(jnp.pi)

print(jax_arr)
print(updated_jax_arr)

# This may seem wasteful, but XLA is smart enough to figure out that
# if one isn't using the input array jax_arr, then it won't allocate
# a special memory object for the output array, it will simply reuse
# the input array and modify it in-place despite not appearing to do so
# at this high-level API.

[[0. 0. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]
[[0.        0.        0.       ]
 [3.1415927 3.1415927 3.1415927]
 [0.        0.        0.       ]]

In [74]:

# However, one still has access to the expressiveness of NumPy!!!
# Example:
print(jax_arr)
another_updated_jax_arr =  jax_arr.at[::2, 1:].add(7)
print(another_updated_jax_arr)

[[0. 0. 0.]
 [0. 0. 0.]
 [0. 0. 0.]]
[[0. 7. 7.]
 [0. 0. 0.]
 [0. 7. 7.]]

$\text{Problem}$: Since JAX wants to be accelerator agnostic, how does JAX handle out-of-bounds indexing?

$\text{Solution}$: Rather than throwing an error message or exception, it simply does a certain “clamping”.

In [78]:

# NumPy behavior

try:
    np.arange(10)[11]
except Exception as e:
    print(f"Exception {e}")

Exception index 11 is out of bounds for axis 0 with size 10

In [ ]:

# JAX behavior is disturbing and likely to be the cause of many bugs!!!
# similar to NaN behavior when doing invalid floating point arithmetic
print(jnp.arange(10).at[11].add(823942))
print(jnp.arange(10)[11])

[0 1 2 3 4 5 6 7 8 9]
9

In [82]:

# Another "gotcha" of JAX:

# NumPy
print(np.sum([1, 2, 3]))

# JAX
try:
    jnp.sum([1, 2, 3])
except Exception as e:
    print(f"TypeError: {e}")

6
TypeError: sum requires ndarray or scalar arguments, got <class 'list'> at position 0.

In [ ]:

# The reason for this behavior can, as with any JAX behavior, be dissected via jaxpr:

def permissive_sum(x):
    return jnp.sum(jnp.array(x))

x = list(range(10)) #[0,..., 9]
print(make_jaxpr(permissive_sum)(x))

# Thus, JAX is good for researchers looking for highly-optimized and flexible programs
# for a beginnner, PyTorch is friendlier

{ lambda ; a:i32[] b:i32[] c:i32[] d:i32[] e:i32[] f:i32[] g:i32[] h:i32[] i:i32[]
    j:i32[]. let
    k:i32[] = convert_element_type[new_dtype=int32 weak_type=False] a
    l:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] k
    m:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b
    n:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] m
    o:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c
    p:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] o
    q:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d
    r:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] q
    s:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e
    t:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] s
    u:i32[] = convert_element_type[new_dtype=int32 weak_type=False] f
    v:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] u
    w:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g
    x:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] w
    y:i32[] = convert_element_type[new_dtype=int32 weak_type=False] h
    z:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] y
    ba:i32[] = convert_element_type[new_dtype=int32 weak_type=False] i
    bb:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] ba
    bc:i32[] = convert_element_type[new_dtype=int32 weak_type=False] j
    bd:i32[1] = broadcast_in_dim[
      broadcast_dimensions=()
      shape=(1,)
      sharding=None
    ] bc
    be:i32[10] = concatenate[dimension=0] l n p r t v x z bb bd
    bf:i32[] = reduce_sum[axes=(0,)] be
  in (bf,) }

In [ ]:

# NumPy - PRNG is stateful!
print(np.random.random()) # PRNG state is advanced
print(np.random.random()) # thus consuming entropy of PRNG
# in both cases, PRNG state change hidden from us...let's expose!
np.random.seed(seed=0) #seed is INITIAL PRNG state!!!

rng_state = np.random.get_state() #get state of PRNG
print(rng_state[2:]) #print this metadata from state

_ = np.random.uniform()
rng_state = np.random.get_state()
print(rng_state[2:])

_ = np.random.uniform()
rng_state = np.random.get_state()
print(rng_state[2:])

# Mersenne Twister PRNG known to have problems (NumPy's implementation of PRNG)

0.6027633760716439
0.5448831829968969
(624, 0, 0.0)
(2, 0, 0.0)
(4, 0, 0.0)

In [ ]:

# In functional programming paradigm, JAX random functions can't modify PRNG state

key = random.PRNGKey(seed=0)
print(key)

print(random.normal(key, shape=(1,)))
print(key)

print(random.normal(key, shape=(1,)))
print(key)

# clearly not random! How to deal with it?

[0 0]
[1.6226422]
[0 0]
[1.6226422]
[0 0]

In [119]:

# Solution: splitting! Not just a key, but key + subkey!

print("Old key:", key)
key, subkey = random.split(key) # can also split into > 1 subkeys
normal_pseudorandom = random.normal(subkey, shape=(1,))
print("Key:", key)
print("Subkey:", subkey)
print("Random number:", normal_pseudorandom)
# functionally, key and subkey are indistinguishable, only a convention
# try running this cell multiple times!

Old key: [ 197075234 2075500836]
Key: [2234728722 1518742019]
Subkey: [3499959921 3652298783]
Random number: [-0.2507795]

In [ ]:

# Why this design?
# Answer: code is reproducible, parallelizable and vectorizable

np.random.seed(seed=0)

def bar(): return np.random.uniform()
def baz(): return np.random.uniform()

def foo(): return bar() + 2 * baz()

print(foo())
# b/c start from same seed state, calling foo() gives same answer each time
# However, if we were to jit foo(), then jit might decide to parallelize the
# program by calling bar() on one core and baz() on another one, hence one may
# obtain different results (e.g. 0.3 + 2*0.4 vs. 0.4 + 2*0.3 which is not reproducible)

1.9791922366721637

In [136]:

# NumPy
np.random.seed(seed=0)
print("Individually:", np.stack([np.random.uniform() for _ in range(3)]))

np.random.seed(seed=0)
print("Simultaneously:", np.random.uniform(size=3))
print("They're the same!")

# JAX
key = random.PRNGKey(seed=0)
subkeys = random.split(key, 3) # creating 3 subkeys
sequence = np.stack([random.normal(subkey) for subkey in subkeys])
print("Individually:", sequence)

key = random.PRNGKey(seed=0)
print("Simultaneously:", random.normal(key, shape=(3,)))
print("They're different!")

Individually: [0.5488135  0.71518937 0.60276338]
Simultaneously: [0.5488135  0.71518937 0.60276338]
They're the same!
Individually: [ 1.0040143 -2.4424558  1.2956359]
Simultaneously: [ 1.6226422   2.0252647  -0.43359444]
They're different!

In [ ]:

# Python control flow + grad() --> no problems!

def f(x):
    if x < 3:
        return 3* x**2
    else:
        return -4*x
    
x = np.linspace(-10, 10, 1000)
y = [f(_) for _ in x]
plt.plot(x, y); plt.show()

print(grad(f)(2.)) #correct!
print(grad(f)(4.)) #correct!

12.0
-4.0

In [147]:

# but if you also want to jit f, then because we're conditioning on x
# in the function scope (if x < 3), it must be passed as a static argument

f_jit = jit(f, static_argnums=(0,))
x = 2

print(make_jaxpr(f_jit, static_argnums=(0,))(x))
print(f_jit(x))

{ lambda ; . let
    a:i32[] = pjit[name=f jaxpr={ lambda ; . let  in (12:i32[],) }] 
  in (a,) }
12

In [156]:

def f(x, n):
    y = 0
    for i in range(n):
        y = y + x[i]
    return y

f_jit = jit(f, static_argnums=(1,))
x = (jnp.array([2, 3, 4]), 15)
print(x)
print(*x)

print(make_jaxpr(f_jit, static_argnums=(1,))(*x))
print(f_jit(*x))
print(2+3+4*13)

(Array([2, 3, 4], dtype=int32), 15)
[2 3 4] 15
{ lambda ; a:i32[3]. let
    b:i32[] = pjit[
      name=f
      jaxpr={ lambda ; a:i32[3]. let
          c:i32[1] = slice[limit_indices=(1,) start_indices=(0,) strides=None] a
          d:i32[] = squeeze[dimensions=(0,)] c
          e:i32[] = add 0:i32[] d
          f:i32[1] = slice[limit_indices=(2,) start_indices=(1,) strides=None] a
          g:i32[] = squeeze[dimensions=(0,)] f
          h:i32[] = add e g
          i:i32[1] = slice[limit_indices=(3,) start_indices=(2,) strides=None] a
          j:i32[] = squeeze[dimensions=(0,)] i
          k:i32[] = add h j
          l:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 3:i32[]
          m:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a l
          n:i32[] = add k m
          o:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 4:i32[]
          p:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a o
          q:i32[] = add n p
          r:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 5:i32[]
          s:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a r
          t:i32[] = add q s
          u:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 6:i32[]
          v:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a u
          w:i32[] = add t v
          x:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 7:i32[]
          y:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a x
          z:i32[] = add w y
          ba:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 8:i32[]
          bb:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a ba
          bc:i32[] = add z bb
          bd:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 9:i32[]
          be:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bd
          bf:i32[] = add bc be
          bg:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 10:i32[]
          bh:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bg
          bi:i32[] = add bf bh
          bj:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 11:i32[]
          bk:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bj
          bl:i32[] = add bi bk
          bm:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 12:i32[]
          bn:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bm
          bo:i32[] = add bl bn
          bp:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 13:i32[]
          bq:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bp
          br:i32[] = add bo bq
          bs:i32[1] = broadcast_in_dim[
            broadcast_dimensions=()
            shape=(1,)
            sharding=None
          ] 14:i32[]
          bt:i32[] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(), collapsed_slice_dims=(0,), start_index_map=(0,), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=True
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(1,)
            unique_indices=True
          ] a bs
          b:i32[] = add br bt
        in (b,) }
    ] a
  in (b,) }
57
57

In [ ]:

# Better way (though less readable) solution is to use low-level Lax API
def f_fori(x, n):
    body_fun = lambda i,val: val+x[i]
    return lax.fori_loop(0, n, body_fun, 0)

f_fori_jit = jit(f_fori)
print(make_jaxpr(f_fori_jit)(*x))
print(f_fori_jit(*x))

{ lambda ; a:i32[3] b:i32[]. let
    c:i32[] = pjit[
      name=f_fori
      jaxpr={ lambda ; a:i32[3] b:i32[]. let
          _:i32[] _:i32[] c:i32[] = while[
            body_jaxpr={ lambda ; d:i32[3] e:i32[] f:i32[] g:i32[]. let
                h:i32[] = add e 1:i32[]
                i:bool[] = lt e 0:i32[]
                j:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e
                k:i32[] = add j 3:i32[]
                l:i32[] = select_n i e k
                m:i32[1] = dynamic_slice[slice_sizes=(1,)] d l
                n:i32[] = squeeze[dimensions=(0,)] m
                o:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g
                p:i32[] = add o n
              in (h, f, p) }
            body_nconsts=1
            cond_jaxpr={ lambda ; q:i32[] r:i32[] s:i32[]. let
                t:bool[] = lt q r
              in (t,) }
            cond_nconsts=0
          ] a 0:i32[] b 0:i32[]
        in (c,) }
    ] a b
  in (c,) }
57

In [ ]:

# Conditioning on data dimensionality is permitted

def log2_if_rank_2(x):
    if x.ndim == 2:
        ln_x = jnp.log(x)
        ln_2 = jnp.log(2)
        return ln_x/ln_2
    else:
        return x
    
print(make_jaxpr(log2_if_rank_2)(jax.numpy.array([1, 2, 3])))
# because array is 3D, just return input as output

{ lambda ; a:i32[3]. let  in (a,) }

In [ ]:

jnp.divide(0, 0)

Out[ ]:

Array(nan, dtype=float32, weak_type=True)

In [160]:

jnp.divide(0, 0)
from jax import config
config.update("jax_debug_nans", True)

Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.

---------------------------------------------------------------------------
FloatingPointError                        Traceback (most recent call last)
Cell In[160], line 1
----> 1 jnp.divide(0, 0)
      2 from jax import config
      3 config.update("jax_debug_nans", True)

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/ufuncs.py:2481, in divide(x1, x2)
   2478 @export
   2479 def divide(x1: ArrayLike, x2: ArrayLike, /) -> Array:
   2480   """Alias of :func:`jax.numpy.true_divide`."""
-> 2481   return true_divide(x1, x2)

    [... skipping hidden 4 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/numpy/ufuncs.py:2473, in true_divide(x1, x2)
   2471 x1, x2 = promote_args_inexact("true_divide", x1, x2)
   2472 jnp_error._set_error_if_divide_by_zero(x2)
-> 2473 out = lax.div(x1, x2)
   2474 jnp_error._set_error_if_nan(out)
   2475 return out

    [... skipping hidden 8 frame]

File ~/jax_linus/lib/python3.10/site-packages/jax/_src/pjit.py:181, in _python_pjit_helper(fun, jit_info, *args, **kwargs)
    179 except api_util.InternalFloatingPointError as e:
    180   if getattr(fun, '_apply_primitive', False):
--> 181     raise FloatingPointError(f"invalid value ({e.ty}) encountered in {fun.__qualname__}") from None
    182   api_util.maybe_recursive_nan_check(e, fun, args, kwargs)
    184 if p.box_data:

FloatingPointError: invalid value (nan) encountered in div

In [165]:

# JAX enforces single precision, though there are ways around this.

x = random.uniform(key, shape=(1000,), dtype=jnp.float64)
print(x.dtype)
print("__________________")
print(type(x))

float32
__________________
<class 'jaxlib._jax.ArrayImpl'>

/tmp/ipykernel_424021/1243311004.py:3: UserWarning: Explicitly requested dtype <class 'jax.numpy.float64'>  is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  x = random.uniform(key, shape=(1000,), dtype=jnp.float64)

Posted in Blog | Leave a comment

Quantumplations

Renormalization Group

Convolutional Neural Networks

Hamilton’s Optics-Mechanics Analogy

Pseudo-Riemannian Geometry

Reinforcement Learning (Part \(1\))

Density Functional Theory

PyTorch Fundamentals (Part \(2\))

PyTorch model building essentials¶

PyTorch Fundamentals (Part \(1\))

Random Tensors¶

Tensor datatypes¶

Manipulating Tensors (Tensor Operation)¶

Self-Attention in Transformers

JAX Fundamentals (Part \(1\))

Archives

Meta