Tensors & Autograd

import torch

# From Python lists
x = torch.tensor([1.0, 2.0, 3.0])
matrix = torch.tensor([[1, 2], [3, 4]])

# Common initialization patterns
zeros = torch.zeros(3, 4)          # 3x4 matrix of zeros
ones = torch.ones(2, 3)            # 2x3 matrix of ones
rand = torch.randn(5, 5)           # 5x5 random normal
eye = torch.eye(3)                 # 3x3 identity matrix
arange = torch.arange(0, 10, 2)   # [0, 2, 4, 6, 8]

# From NumPy (zero-copy when possible)
import numpy as np
np_array = np.array([1, 2, 3])
tensor_from_np = torch.from_numpy(np_array)
back_to_np = tensor_from_np.numpy()

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

# Arithmetic (element-wise)
print(a + b)        # or torch.add(a, b)
print(a * b)        # Element-wise multiply

# Matrix operations
print(a @ b)        # Matrix multiplication
print(a.T)          # Transpose

# Reductions
print(a.mean())     # Mean of all elements
print(a.sum(dim=0)) # Sum along rows
print(a.max())      # Maximum value

# Reshaping
x = torch.arange(12)
print(x.view(3, 4))     # Reshape to 3x4
print(x.reshape(2, 6))  # Same, but works with non-contiguous
print(x.unsqueeze(0))   # Add batch dimension

# Check GPU availability
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using: {device}")

# Move tensors to GPU
x = torch.randn(1000, 1000).to(device)
y = torch.randn(1000, 1000).to(device)

# Operations on GPU tensors happen on GPU
z = x @ y  # This runs on GPU!

# Move back to CPU when needed
result = z.cpu().numpy()

# Simple gradient computation
x = torch.tensor([2.0, 3.0], requires_grad=True)

# Forward pass: y = x^2 + 3x
y = x ** 2 + 3 * x
loss = y.sum()

# Backward pass: compute gradients
loss.backward()

# dy/dx = 2x + 3
print(x.grad)  # tensor([7., 9.]) because 2*2+3=7, 2*3+3=9

# Important: zero gradients before next backward pass!
x.grad.zero_()

# Disable gradient tracking (for inference)
with torch.no_grad():
    predictions = model(test_data)  # No graph built, faster

# Linear regression from scratch with autograd
import torch

# Data: y = 2x + 1
x = torch.linspace(0, 10, 100)
y = 2 * x + 1 + torch.randn(100) * 0.5

# Learnable parameters
w = torch.tensor([0.0], requires_grad=True)
b = torch.tensor([0.0], requires_grad=True)

lr = 0.001
for epoch in range(100):
    # Forward
    y_pred = w * x + b
    loss = ((y_pred - y) ** 2).mean()

    # Backward
    loss.backward()

    # Update weights (no gradient tracking needed)
    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad

    # Zero gradients
    w.grad.zero_()
    b.grad.zero_()

print(f"w={w.item():.2f}, b={b.item():.2f}")  # Close to w=2, b=1