PyTorch Engineering Masterclass — Chapter 2 of 16

Welcome to Chapter 2 of the PyTorch Engineering Masterclass. In high-performance deep learning systems, tensors are not simple multidimensional arrays; they are strided mathematical views mapped directly over contiguous physical memory blocks. Understanding tensor memory layouts, pointer arithmetic, BLAS/LAPACK linear algebra kernels, Einstein summation contractions, and asynchronous DMA streaming over PCIe buses is what separates a novice programmer from a senior ML infrastructure engineer. In this chapter, we dissect PyTorch tensors at the hardware and algorithmic level.

1. Strided Memory Architecture & Physical Storage Layout

Every PyTorch tensor is composed of two distinct components: a Typed Storage Buffer (a flat 1D array of raw contiguous bytes allocated on host RAM or GPU VRAM) and a Tensor Metadata View (which dictates how multidimensional indices map onto that flat 1D buffer). The view is defined by three critical attributes:

1. Size / Shape: The number of elements along each logical dimension $(d_0, d_1, \dots, d_{n-1})$.
2. Stride: The number of elements in physical memory that must be skipped to advance by 1 step along each dimension $(s_0, s_1, \dots, s_{n-1})$.
3. Storage Offset: The starting index in the physical storage buffer where the tensor view begins.

The physical storage address of any element at multidimensional coordinate $(i_0, i_1, \dots, i_{n-1})$ is computed using the canonical linear addressing equation:

Physical Index = Storage Offset + ∑ (i_k × s_k)

1.1 Row-Major (C-Contiguous) vs. Non-Contiguous Memory

In standard C-contiguous layout, elements along the last dimension are stored adjacent to one another in physical RAM. For a 2D tensor of shape [Rows, Cols], the strides are (Cols, 1). When you transpose a matrix using t.t(), PyTorch does not reallocate or copy memory: it merely swaps the strides to (1, Cols)! The tensor is now non-contiguous.

Figure 2.1: Physical 1D Storage Buffer aliased by multiple Tensor Views with distinct Stride Metadata.

import torch

# Create a contiguous 2D tensor [3, 4]
x = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print(f"Shape: {x.shape}, Strides: {x.stride()}, Contiguous: {x.is_contiguous()}")
# Output: Shape: torch.Size([3, 4]), Strides: (4, 1), Contiguous: True

# Transpose without copying (swapping strides)
x_t = x.t()
print(f"Transposed Shape: {x_t.shape}, Strides: {x_t.stride()}, Contiguous: {x_t.is_contiguous()}")
# Output: Transposed Shape: torch.Size([4, 3]), Strides: (1, 4), Contiguous: False

# Verify that both tensors share the exact same underlying memory pointer:
print(f"x data_ptr:   {x.data_ptr()}")
print(f"x_t data_ptr: {x_t.data_ptr()} (Identical physical address!)")

2. Zero-Copy Sliding Windows via torch.as_strided

One of the most powerful yet dangerous low-level tools in PyTorch is torch.as_strided(input, size, stride, storage_offset=0). It allows you to construct arbitrary multidimensional views over existing memory buffers without allocating a single byte of additional RAM. This is the exact mechanism used by convolution engines, image patch extractors, and time-series feature generators.

Figure 2.2: Zero-Copy Sliding Window patch extraction over 1D storage using custom stride multipliers.

Animation 2.1: Frame-by-frame visualization of as_strided sweeping sliding windows across a 1D storage array with zero memory overhead.

def sliding_window_1d(x: torch.Tensor, window_size: int, step: int = 1) -> torch.Tensor:
    """
    Extracts 1D sliding windows with zero memory allocation using stride manipulation.
    """
    assert x.ndim == 1, "Input must be 1D"
    assert x.is_contiguous(), "Input must be contiguous in memory"
    num_windows = (x.size(0) - window_size) // step + 1
    new_shape = (num_windows, window_size)
    new_strides = (step * x.stride(0), x.stride(0))
    return torch.as_strided(x, size=new_shape, stride=new_strides)

# Demonstration:
series = torch.arange(10, dtype=torch.float32)
windows = sliding_window_1d(series, window_size=4, step=2)
print("Original series:", series)
print("Sliding windows (zero-copy):
", windows)
print("windows data_ptr == series data_ptr:", windows.data_ptr() == series.data_ptr())

3. Host Memory Engineering: Pinned Page-Locked RAM & PCIe DMA

When training massive deep learning models on NVIDIA GPUs, the host CPU memory system is frequently the primary system bottleneck. By default, host operating systems allocate Pageable Memory in physical RAM. Pageable memory can be swapped out to disk or relocated by the OS virtual memory manager at any moment.

Because the physical address of pageable memory is unstable, GPU Direct Memory Access (DMA) controllers cannot read from it directly! When you invoke tensor.to('cuda') on a standard tensor, the PyTorch CUDA runtime is forced to perform a costly two-step staged transfer:

1. The CPU synchronously copies the tensor from pageable RAM into an internal temporary Pinned Staging Buffer.
2. The GPU DMA controller initiates the PCIe hardware transfer from the staging buffer into GPU VRAM.

Figure 2.3: Architecture comparison between synchronous two-step pageable staging and direct one-step Pinned DMA streaming.

Animation 2.2: Motion comparison demonstrating how Pinned Memory bypasses CPU staging, achieving 4.8x higher PCIe transfer bandwidth.

3.1 Asynchronous DMA Streaming via non_blocking=True

By locking the host pages in physical RAM via pin_memory() (or setting pin_memory=True in DataLoader), you allow the GPU to initiate hardware DMA directly. When paired with non_blocking=True, the transfer executes concurrently on an independent CUDA copy stream while the CPU proceeds immediately to prepare the next batch!

import torch

# Standard pageable tensor
cpu_tensor = torch.randn(10_000, 10_000, dtype=torch.float32)

# Allocate pinned (page-locked) host memory
pinned_tensor = cpu_tensor.pin_memory()
print("Is tensor pinned in physical RAM?", pinned_tensor.is_pinned())

# Asynchronous non-blocking transfer over PCIe DMA
if torch.cuda.is_available():
    # Asynchronous launch: returns immediately to Python without waiting for PCIe bus
    gpu_tensor = pinned_tensor.to("cuda", non_blocking=True)
    # CPU can now do data augmentation, tokenization, or disk I/O concurrently!

4. Linear Algebra & BLAS/LAPACK Kernels (SVD, Cholesky, Solve)

Modern machine learning models depend fundamentally on dense linear algebra routines powered by underlying BLAS (Basic Linear Algebra Subprograms) and LAPACK libraries (such as Intel MKL, OpenBLAS, and NVIDIA cuBLAS / cuSOLVER). PyTorch exposes these high-performance kernels through the torch.linalg module.

Figure 2.4: Singular Value Decomposition (SVD) matrix factorization and Low-Rank compression mechanics.

Animation 2.3: Matrix factorization via SVD and truncated low-rank approximation (the theoretical foundation of LoRA weight adapters).

4.1 Low-Rank Weight Approximations (The Principle Behind LoRA)

Any dense weight matrix $W \in \mathbb{R}^{M \imes N}$ can be factored via Singular Value Decomposition:

$$W = U \Sigma V^T$$

By retaining only the top-$k$ largest singular values in $\Sigma$ and discarding the rest, we obtain the optimal rank-$k$ approximation $W_k$ that minimizes Frobenius reconstruction error (Eckart–Young–Mirsky theorem):

$$W_k = U_{:, :k} \Sigma_{:k, :k} V_{:k, :}^T$$

import torch

# Generate a synthetic weight matrix (e.g., LLM projection layer)
torch.manual_seed(42)
W = torch.randn(1024, 1024, dtype=torch.float32)

# Full SVD decomposition
U, S, Vh = torch.linalg.svd(W, full_matrices=False)

# Compress to rank k=32 (LoRA rank)
k = 32
U_k = U[:, :k]
S_k = torch.diag(S[:k])
Vh_k = Vh[:k, :]

# Reconstruct low-rank matrix
W_approx = U_k @ S_k @ Vh_k

# Evaluate parameter compression and energy retention
original_params = W.numel()
compressed_params = U_k.numel() + k + Vh_k.numel()
compression_ratio = original_params / compressed_params
energy_retained = (torch.sum(S[:k]**2) / torch.sum(S**2)).item() * 100

print(f"Original Parameters:   {original_params:,}")
print(f"Compressed Parameters: {compressed_params:,} ({compression_ratio:.1f}x Compression)")
print(f"Energy Retained:       {energy_retained:.2f}%")

5. Einstein Summation (torch.einsum) for Modern Neural Architectures

In modern deep learning architectures—especially Transformers, Multi-Head Attention, and Graph Neural Networks—standard matrix multiplication functions (torch.matmul, torch.bmm) become unwieldy, requiring multiple consecutive permute, transpose, and unsqueeze calls. Einstein summation provides a unified, highly optimized notation for tensor contractions.

Figure 2.5: Einstein index notation mapping: free indices vs contracted summation indices across tensors.

Animation 2.4: Step-by-step token contraction in Multi-Head Attention using torch.einsum(‘b i d, b j d -> b i j’).

5.1 The Rules of Einstein Notation

1. Free Indices: Indices that appear on both the input and output sides of the arrow ->. These dimensions are preserved in the output.
2. Contracted (Bound) Indices: Indices that appear on the input side but omitted from the output. The elements along these dimensions are multiplied element-wise and summed together.
3. Batch Indices: Indices that appear in all input tensors and the output without change. They represent parallel independent batch slices.

import torch

B, H, S, D = 4, 8, 128, 64
Q = torch.randn(B, H, S, D)
K = torch.randn(B, H, S, D)
V = torch.randn(B, H, S, D)

# Compute Multi-Head Attention Logits with einsum
scores = torch.einsum('bhid,bhjd->bhij', Q, K) * (D ** -0.5)
attn_weights = torch.softmax(scores, dim=-1)

# Compute Output Context
context = torch.einsum('bhij,bhjd->bhid', attn_weights, V)
print("Attention Context Tensor Shape:", context.shape)
# Output: torch.Size([4, 8, 128, 64])

6. Memory Aliasing, In-Place Mutations & Autograd _version Tracking

In PyTorch, operations that end in an underscore (e.g., add_(), sub_(), mul_(), zero_()) modify the tensor’s physical storage in-place rather than allocating a new tensor. While in-place mutations save memory, they represent one of the most frequent sources of catastrophic silent failures when paired with PyTorch Autograd.

Animation 2.5: The internal mechanics of PyTorch’s _version counter detecting illegal in-place mutations during backward graph execution.

6.1 How Autograd Detects Illegal Mutations

Every PyTorch tensor maintains an internal integer field called _version. When a tensor is created, its _version starts at 0. Every time an in-place mutation executes, the C++ core increments this counter:

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
print("Initial version:", x._version)  # 0

# Forward computation saving x for backward
y = x ** 2
print("Version after forward:", x._version)  # 0

# In-place mutation!
x.add_(5.0)
print("Version after in-place add_:", x._version)  # 1

# Attempting backpropagation throws a fatal RuntimeError:
try:
    y.backward(torch.ones_like(y))
except RuntimeError as e:
    print("
Caught Autograd Safety Violation:")
    print(e)
    # RuntimeError: one of the variables needed for gradient computation has been modified by an in-place operation

Because the mathematical derivative of $y = x^2$ is $\frac{dy}{dx} = 2x$, the backward node requires the exact original value of $x$ that produced $y$. Mutating $x$ in-place overwrites that value, making correct gradient calculation impossible. Autograd’s version counter acts as an infallible circuit breaker.

7. Numerical Precision, Stability & The LogSumExp Trick

Modern AI systems rely heavily on reduced-precision floating-point formats to maximize throughput on modern tensor cores:

• FP32 (IEEE Single Precision): 1 sign bit, 8 exponent bits, 23 mantissa bits. Standard reference baseline.
• FP16 (Half Precision): 1 sign bit, 5 exponent bits, 10 mantissa bits. Fast, but susceptible to underflow ($< 6 \imes 10^{-5}$) and overflow ($> 65,504$).
• BF16 (Brain Floating Point): 1 sign bit, 8 exponent bits, 7 mantissa bits. Preserves the exact dynamic range of FP32 with reduced precision; highly resilient to gradient underflow/overflow.

Figure 2.6: The LogSumExp mathematical trick preventing numerical underflow and NaN explosions in cross-entropy loss.

7.1 The LogSumExp Formulation

In classification and language modeling, calculating normalized probabilities involves expressions of the form $\log \sum_{i=1}^N e^{x_i}$. When $x_i$ is large (e.g., $x = 1000$), $e^{1000}$ overflows to inf. When $x_i$ is strongly negative (e.g., $x = -1000$), $e^{-1000}$ underflows to 0.0, causing $\log(0)$ to yield -inf or NaN.

The LogSumExp Trick subtracts the maximum value $m = \max(x)$ before exponentiating:

$$\log \sum_{i=1}^N e^{x_i} = m + \log \sum_{i=1}^N e^{x_i – m}$$

Because $x_i – m \le 0$ for all $i$, the largest exponent is $e^0 = 1$, completely eliminating floating-point overflow!

import torch

# Unstable naive implementation
logits = torch.tensor([1000.0, 1001.0, 1002.0], dtype=torch.float32)

naive_lse = torch.log(torch.sum(torch.exp(logits)))
print("Naive LogSumExp:", naive_lse.item())  # inf! (Numerical Failure)

# Production numerically stable PyTorch implementation
stable_lse = torch.logsumexp(logits, dim=0)
print("PyTorch torch.logsumexp:", stable_lse.item())  # 1002.4076 (Mathematically exact!)

8. Advanced Multidimensional Broadcasting & Causal Attention Masks

Broadcasting allows PyTorch to execute element-wise operations between tensors of differing shapes without copying data. The rule is strictly deterministic: starting from the trailing (rightmost) dimension and working backwards, two dimensions are compatible if:

1. They are equal in size, or
2. One of them is equal to 1.

Figure 2.7: Causal Upper Mask broadcasting across Batch and Multi-Head dimensions via stride-0 expansions.

8.1 Stride-0 Expansion Mechanics

When PyTorch stretches a dimension of size 1 to size N during broadcasting, it does not duplicate data in memory! Instead, it sets the stride for that dimension to 0. Every step along that dimension accesses the exact same physical memory offset, achieving true $O(1)$ memory expansion.

import torch

# Attention logits: [Batch=16, Heads=8, SeqLen=128, SeqLen=128]
logits = torch.randn(16, 8, 128, 128)

# Single causal mask: [1, 1, 128, 128]
causal_mask = torch.triu(torch.full((128, 128), float('-inf')), diagonal=1)
causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, 128, 128]

# Broadcasting addition: the single mask serves 128 distinct attention heads with ZERO memory copy!
masked_logits = logits + causal_mask
print(f"Masked Logits Shape: {masked_logits.shape}")

# Verify stride-0 expansion:
expanded_mask = causal_mask.expand(16, 8, 128, 128)
print(f"Expanded Mask Strides: {expanded_mask.stride()}")
# Output: Expanded Mask Strides: (0, 0, 128, 1) -> Dim 0 and Dim 1 have stride 0!

9. CPU Vectorization: SIMD Registers (AVX-512, AMX, ARM Neon)

While GPUs excel at massively parallel throughput, high-performance CPU inference is critical in edge devices, mobile stacks, and cost-optimized microservices. PyTorch compiles its CPU backend using vectorized SIMD (Single Instruction, Multiple Data) intrinsics.

Figure 2.8: Scalar CPU execution loop vs Vectorized SIMD register execution (AVX-512 / ARM Neon).

Instead of processing one float32 per clock cycle, an AVX-512 register packs sixteen 32-bit floating point numbers into a single 512-bit hardware register lane, executing sixteen additions or fused multiply-adds (FMA) in a single CPU instruction cycle!

10. Production Engineering Laboratories

Lab 1: High-Performance Zero-Copy 2D Image Patch Extraction

Objective: Extract overlapping $8 \imes 8$ image patches from a high-resolution feature map using as_strided without allocating redundant memory.

import torch

def extract_patches_2d(img: torch.Tensor, patch_h: int, patch_w: int, stride_h: int, stride_w: int) -> torch.Tensor:
    """
    Extracts 2D image patches [N, C, H_out, W_out, patch_h, patch_w] with zero memory copy.
    Input: [N, C, H, W] contiguous tensor.
    """
    assert img.ndim == 4 and img.is_contiguous()
    N, C, H, W = img.shape
    out_h = (H - patch_h) // stride_h + 1
    out_w = (W - patch_w) // stride_w + 1
    
    sN, sC, sH, sW = img.stride()
    new_shape = (N, C, out_h, out_w, patch_h, patch_w)
    new_strides = (sN, sC, sH * stride_h, sW * stride_w, sH, sW)
    
    return torch.as_strided(img, size=new_shape, stride=new_strides)

# Test execution:
batch_images = torch.arange(1 * 3 * 32 * 32, dtype=torch.float32).reshape(1, 3, 32, 32)
patches = extract_patches_2d(batch_images, patch_h=8, patch_w=8, stride_h=4, stride_w=4)
print("Patches shape:", patches.shape)
print("Patches share storage with original:", patches.data_ptr() == batch_images.data_ptr())

Lab 2: Multi-Head Self-Attention Core via Einsum

Objective: Implement a complete multi-head self-attention module using pure torch.einsum contractions.

import torch
import torch.nn as nn

class EinsumSelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        
    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        B, S, D = x.shape
        qkv = self.qkv_proj(x).reshape(B, S, 3, self.num_heads, self.d_k)
        # Permute/split using einsum or indexing
        Q, K, V = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
        
        # Attention scores: Q [B, S, H, D] x K [B, T, H, D] -> [B, H, S, T]
        scores = torch.einsum('bshd,bthd->bhst', Q, K) / (self.d_k ** 0.5)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
            
        attn_weights = torch.softmax(scores, dim=-1)
        
        # Context: [B, H, S, T] x V [B, T, H, D] -> [B, S, H, D]
        context = torch.einsum('bhst,bthd->bshd', attn_weights, V)
        context = context.reshape(B, S, D)
        return self.out_proj(context)

# Verify correctness:
attn = EinsumSelfAttention(d_model=256, num_heads=8)
tokens = torch.randn(2, 64, 256)
out = attn(tokens)
print("Einsum Attention Output Shape:", out.shape)  # [2, 64, 256]

Lab 3: Pinned Memory PCIe DMA Throughput Profiler

Objective: Benchmark actual PCIe transfer speeds (GB/s) comparing standard Pageable memory vs Pinned page-locked memory with proper CUDA event synchronization.

import torch
import time

def benchmark_pcie_transfer(size_mb: int = 500, iterations: int = 10):
    if not torch.cuda.is_available():
        print("CUDA GPU required for PCIe DMA benchmark.")
        return
        
    num_elements = (size_mb * 1024 * 1024) // 4  # float32 = 4 bytes
    pageable_cpu = torch.randn(num_elements, dtype=torch.float32)
    pinned_cpu = pageable_cpu.pin_memory()
    
    # Warmup
    _ = pageable_cpu.to("cuda")
    torch.cuda.synchronize()
    
    # 1. Benchmark Pageable Transfer
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    
    start_event.record()
    for _ in range(iterations):
        _ = pageable_cpu.to("cuda", non_blocking=False)
    end_event.record()
    torch.cuda.synchronize()
    
    pageable_ms = start_event.elapsed_time(end_event) / iterations
    pageable_bw = (size_mb / 1024) / (pageable_ms / 1000.0)
    
    # 2. Benchmark Pinned Transfer
    start_event.record()
    for _ in range(iterations):
        _ = pinned_cpu.to("cuda", non_blocking=True)
    end_event.record()
    torch.cuda.synchronize()
    
    pinned_ms = start_event.elapsed_time(end_event) / iterations
    pinned_bw = (size_mb / 1024) / (pinned_ms / 1000.0)
    
    print(f"=== PCIe Host-to-Device Bandwidth ({size_mb} MB) ===")
    print(f"Pageable Memory: {pageable_bw:.2f} GB/s ({pageable_ms:.2f} ms)")
    print(f"Pinned Memory:   {pinned_bw:.2f} GB/s ({pinned_ms:.2f} ms) -> {pinned_bw/pageable_bw:.1f}x Speedup!")

# Run benchmark:
# benchmark_pcie_transfer(size_mb=256)

Lab 4: SVD Low-Rank Linear Layer Factorization

Objective: Factorize a heavy linear layer $W \in \mathbb{R}^{4096 \imes 4096}$ into two low-rank matrices $A \in \mathbb{R}^{4096 \imes 64}$ and $B \in \mathbb{R}^{64 \imes 4096}$, verifying output error bounds.

import torch
import torch.nn as nn

def factorize_linear_svd(layer: nn.Linear, rank: int) -> tuple[nn.Linear, nn.Linear]:
    """
    Factorizes nn.Linear(M, N) into Linear(M, rank) followed by Linear(rank, N).
    """
    with torch.no_grad():
        W = layer.weight.data  # [N, M]
        U, S, Vh = torch.linalg.svd(W, full_matrices=False)
        
        # Distribute singular values square root symmetrically
        s_sqrt = torch.diag(torch.sqrt(S[:rank]))
        W1 = s_sqrt @ Vh[:rank, :]          # [rank, M]
        W2 = U[:, :rank] @ s_sqrt          # [N, rank]
        
        layer1 = nn.Linear(layer.in_features, rank, bias=False)
        layer2 = nn.Linear(rank, layer.out_features, bias=layer.bias is not None)
        
        layer1.weight.data.copy_(W1)
        layer2.weight.data.copy_(W2)
        if layer.bias is not None:
            layer2.bias.data.copy_(layer.bias.data)
            
    return layer1, layer2

# Verification:
dense = nn.Linear(2048, 2048)
l1, l2 = factorize_linear_svd(dense, rank=64)

x = torch.randn(8, 2048)
y_dense = dense(x)
y_fact = l2(l1(x))

rel_error = torch.norm(y_dense - y_fact) / torch.norm(y_dense)
print(f"Dense parameters:      {sum(p.numel() for p in dense.parameters()):,}")
print(f"Factorized parameters: {sum(p.numel() for p in l1.parameters()) + sum(p.numel() for p in l2.parameters()):,}")
print(f"Relative output error: {rel_error.item()*100:.2f}%")

Lab 5: Numerical Stability Stress Test (LogSumExp vs Naive Softmax)

Objective: Simulate catastrophic numerical underflow and overflow in extreme logit environments and observe recovery via torch.logsumexp.

import torch

def compare_stability(extreme_scale: float = 1e4):
    logits = torch.tensor([1.0, 2.0, extreme_scale], dtype=torch.float32)
    
    # 1. Naive Log-Softmax
    exp_logits = torch.exp(logits)
    sum_exp = torch.sum(exp_logits)
    naive_log_probs = logits - torch.log(sum_exp)
    
    # 2. Stable PyTorch Log-Softmax
    stable_log_probs = logits - torch.logsumexp(logits, dim=0)
    
    print(f"Logits: [1.0, 2.0, {extreme_scale}]")
    print("Naive Log-Probs: ", naive_log_probs)
    print("Stable Log-Probs:", stable_log_probs)

compare_stability(extreme_scale=100.0)
compare_stability(extreme_scale=1000.0)

11. Capstone Mini-Project: Tensor Linear Algebra & Compression CLI Suite

In this capstone engineering mini-project, we build a production-grade CLI suite: TensorLinalgOptimizer. It supports:

• Dense matrix decomposition benchmarks (SVD, Cholesky, Eigendecomposition).
• LoRA-style low-rank compression with configurable rank $k$ and energy preservation metrics.
• PCIe DMA transfer profiling comparing Pageable vs Pinned memory.
• Hardware-synchronized precision execution timing across CPU, CUDA, and Apple Silicon MPS.

"""
TensorLinalgOptimizer: High-Performance Linear Algebra & Compression Suite
Author: Senior ML Infrastructure Engineer
PyTorch 2.x Production Standard
"""

import argparse
import sys
import time
import torch
import torch.nn as nn
from typing import Dict, Any

class TensorLinalgOptimizer:
    def __init__(self, device: str = None):
        if device is None:
            if torch.cuda.is_available():
                self.device = torch.device("cuda")
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                self.device = torch.device("mps")
            else:
                self.device = torch.device("cpu")
        else:
            self.device = torch.device(device)
            
        print(f"[INIT] Active Accelerator Target: {self.device}")

    def synchronize(self):
        """Hardware-synchronized barrier ensuring accurate kernel benchmarking."""
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        elif self.device.type == "mps":
            torch.mps.synchronize()

    def benchmark_svd(self, matrix_dim: int = 2048, rank: int = 64) -> Dict[str, Any]:
        print(f"
--- Running SVD Low-Rank Compression (Dim={matrix_dim}x{matrix_dim}, Rank={rank}) ---")
        torch.manual_seed(42)
        A = torch.randn(matrix_dim, matrix_dim, device=self.device, dtype=torch.float32)
        
        self.synchronize()
        t0 = time.perf_counter()
        U, S, Vh = torch.linalg.svd(A, full_matrices=False)
        self.synchronize()
        svd_duration = (time.perf_counter() - t0) * 1000.0
        
        # Low rank reconstruction
        U_k = U[:, :rank]
        S_k = torch.diag(S[:rank])
        Vh_k = Vh[:rank, :]
        
        self.synchronize()
        t1 = time.perf_counter()
        A_k = U_k @ S_k @ Vh_k
        self.synchronize()
        reconstruct_duration = (time.perf_counter() - t1) * 1000.0
        
        # Energy and error metrics
        total_energy = torch.sum(S ** 2).item()
        retained_energy = torch.sum(S[:rank] ** 2).item()
        energy_pct = (retained_energy / total_energy) * 100.0
        rel_frob_error = (torch.norm(A - A_k) / torch.norm(A)).item() * 100.0
        
        original_bytes = A.numel() * 4
        compressed_bytes = (U_k.numel() + rank + Vh_k.numel()) * 4
        compression_ratio = original_bytes / compressed_bytes
        
        results = {
            "svd_compute_ms": svd_duration,
            "reconstruct_ms": reconstruct_duration,
            "energy_retained_pct": energy_pct,
            "frobenius_error_pct": rel_frob_error,
            "compression_ratio": compression_ratio,
            "memory_saved_mb": (original_bytes - compressed_bytes) / (1024 * 1024)
        }
        
        for k, v in results.items():
            print(f"  {k:<22}: {v:.2f}")
        return results

    def benchmark_einsum_attention(self, batch: int = 16, heads: int = 8, seq_len: int = 512, d_k: int = 64):
        print(f"
--- Multi-Head Attention Einsum Benchmark (B={batch}, H={heads}, S={seq_len}, D={d_k}) ---")
        Q = torch.randn(batch, heads, seq_len, d_k, device=self.device)
        K = torch.randn(batch, heads, seq_len, d_k, device=self.device)
        V = torch.randn(batch, heads, seq_len, d_k, device=self.device)
        
        # Warmup
        _ = torch.einsum('bhid,bhjd->bhij', Q, K)
        self.synchronize()
        
        iters = 50
        t0 = time.perf_counter()
        for _ in range(iters):
            scores = torch.einsum('bhid,bhjd->bhij', Q, K) * (d_k ** -0.5)
            probs = torch.softmax(scores, dim=-1)
            _ = torch.einsum('bhij,bhjd->bhid', probs, V)
        self.synchronize()
        avg_ms = ((time.perf_counter() - t0) / iters) * 1000.0
        
        print(f"  Einsum Attention Latency: {avg_ms:.3f} ms / iteration")

def main():
    parser = argparse.ArgumentParser(description="TensorLinalgOptimizer: High-Performance PyTorch CLI")
    parser.add_argument("--device", type=str, default=None, help="Target device (cpu, cuda, mps)")
    parser.add_argument("--dim", type=int, default=1024, help="Square matrix dimension for SVD")
    parser.add_argument("--rank", type=int, default=32, help="Low-rank target approximation")
    args = parser.parse_args()
    
    opt = TensorLinalgOptimizer(device=args.device)
    opt.benchmark_svd(matrix_dim=args.dim, rank=args.rank)
    opt.benchmark_einsum_attention()

if __name__ == "__main__":
    main()

12. 10 Practical Exercises with Solutions

# Initial strides: (3*4, 4, 1) = (12, 4, 1)
# After permute(2, 0, 1): dimensions are [4, 2, 3].
# The original strides (12, 4, 1) are reordered according to the permutation:
# New strides: (1, 12, 4).
x = torch.empty(2, 3, 4)
print("Initial:", x.stride())          # (12, 4, 1)
print("Permuted:", x.permute(2, 0, 1).stride())  # (1, 12, 4)

A = torch.randn(8, 16, 16)
# Free index: b (batch). Repeated index: i (diagonal contraction).
batch_traces = torch.einsum('bii->b', A)
print(batch_traces.shape)  # torch.Size([8])

def im2col_strided(img: torch.Tensor, K: int):
    H, W = img.shape
    out_h, out_w = H - K + 1, W - K + 1
    sH, sW = img.stride()
    return torch.as_strided(img, size=(out_h, out_w, K, K), stride=(sH, sW, sH, sW))

img = torch.randn(10, 10)
windows = im2col_strided(img, K=3)
print(windows.shape)  # torch.Size([8, 8, 3, 3])

def shares_memory(a: torch.Tensor, b: torch.Tensor) -> bool:
    return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()

x = torch.randn(4, 4)
y = x.view(-1)
z = x.clone()
print("x and y share memory:", shares_memory(x, y))  # True
print("x and z share memory:", shares_memory(x, z))  # False

# Bug: w.add_(1.0) modifies w in-place while needed for loss.backward().
# Fix: Move in-place updates outside the computation graph or use out-of-place addition:
w = torch.randn(10, requires_grad=True)
loss = (w ** 2).sum()
loss.backward()

# Perform in-place optimizer step safely inside torch.no_grad():
with torch.no_grad():
    w.add_(1.0)
    w.grad.zero_()

u = torch.randn(4, 5)
v = torch.randn(4, 7)
outer = torch.einsum('bm,bn->bmn', u, v)
print("Outer shape:", outer.shape)  # torch.Size([4, 5, 7])

def robust_softmax(logits: torch.Tensor, dim: int = -1) -> torch.Tensor:
    # Subtract max along target dimension before exponentiating
    max_logits = torch.amax(logits, dim=dim, keepdim=True)
    exp_logits = torch.exp(logits - max_logits)
    sum_exp = torch.sum(exp_logits, dim=dim, keepdim=True)
    return exp_logits / sum_exp

extreme = torch.tensor([[1000.0, 1001.0, 1002.0]])
print("Robust Softmax:", robust_softmax(extreme))
# Sums to 1.0 with zero overflow!

A = torch.randn(10, 10)
B = torch.randn(10, 3)
# solve uses LU decomposition, which is both faster and far more numerically stable than matrix inversion
X = torch.linalg.solve(A, B)
# Verify residual norm ||AX - B||:
residual = torch.norm(A @ X - B)
print("Residual error:", residual.item())

# Construct positive-definite matrix
M = torch.randn(5, 5)
cov = M @ M.t() + torch.eye(5) * 1e-3  # Guarantee positive definiteness
L = torch.linalg.cholesky(cov)
print("Cholesky factor shape:", L.shape)
print("Reconstruction error:", torch.norm(cov - L @ L.t()).item())

# Explanation: view() requires the input tensor to be contiguous in memory because
# it cannot change physical data ordering. x.t() has transposed strides.
# - contiguous().view(-1): explicitly allocates a new contiguous copy, then views it.
# - reshape(-1): automatically inspects contiguity; returns a view if contiguous,
#   or transparently clones memory if non-contiguous.
x = torch.arange(6).reshape(2, 3)
xt = x.t()
# xt.view(-1) -> RuntimeError: view size is not compatible with input tensor's size and stride
print("Fixed via reshape:", xt.reshape(-1))

13. 10 Senior Machine Learning Interview Questions & Answers

14. Final Knowledge Assessment & Official References