import sys
import torch
import torchvision

print("Python executable:", sys.executable)
print("PyTorch version:", torch.__version__)
print("Torchvision version:", torchvision.__version__)

import copy
import math

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import datasets, transforms

import a1_v8_support as support
from a1_v8_checks import check, check_indexing_practice, check_broadcasting_practice

support.seed_all(seed=0)
device = torch.device("cpu")
results = {}

toy_scores = torch.tensor([[2.0, 1.0, 0.0], [0.0, 3.0, 1.0]], dtype=torch.float32)
toy_labels = torch.tensor([0, 2], dtype=torch.long)
class_bias = torch.tensor([0.1, 0.2, 0.3])

print(toy_scores)
print("Dimensions:", toy_scores.ndim, "| shape:", tuple(toy_scores.shape))
print("Number type:", toy_scores.dtype, "| stored on:", toy_scores.device)

def indexing_practice(scores, labels):
    first_row = None  # TODO: first example, keeping shape (1,C).
    last_two_columns = None  # TODO: last two class columns, shape (B,2).
    class_zero_rows = None  # TODO: examples whose label is 0, shape (K,C).
    target_scores = None  # TODO: each example's correct-class score, shape (B,).
    raise NotImplementedError("Complete the four expressions, then remove this line.")
    return first_row, last_two_columns, class_zero_rows, target_scores

check_indexing_practice(indexing_practice)
first_row, last_columns, zero_rows, target_scores = indexing_practice(
    toy_scores, toy_labels
)
print("First example with batch axis:\n", first_row)
print("Last two columns:\n", last_columns)
print("Examples labeled 0:\n", zero_rows)
print("Correct-class scores:", target_scores)

def broadcasting_practice(scores, bias):
    biased_scores = None  # TODO: add the class biases, shape (B,C).
    row_maximum = None  # TODO: one maximum per row, shape (B,1).
    shifted_scores = None  # TODO: subtract from the original scores, shape (B,C).
    raise NotImplementedError("Complete the three expressions, then remove this line.")
    return biased_scores, shifted_scores

check_broadcasting_practice(broadcasting_practice)
biased_scores, shifted_scores = broadcasting_practice(toy_scores, class_bias)
print("With class biases:\n", biased_scores, "| shape:", tuple(biased_scores.shape))
print(
    "After subtracting row maxima:\n",
    shifted_scores,
    "| shape:",
    tuple(shifted_scores.shape),
)

def build_datasets(root, split):
    transform = None  # TODO: create the image-to-tensor transform.
    mnist_train = None  # TODO: create the official training Dataset.
    mnist_test = None  # TODO: create the official test Dataset.

    train_dataset = None  # TODO: Subset of mnist_train.
    validation_dataset = None  # TODO: another Subset of mnist_train.
    test_dataset = None  # TODO: Subset of mnist_test.
    raise NotImplementedError("Create the datasets, then remove this line.")
    return train_dataset, validation_dataset, test_dataset

with np.load(support.DATA / "mnist_splits.npz") as split_file:
    split = {key: split_file[key].copy() for key in split_file.files}

train_dataset, validation_dataset, test_dataset = build_datasets(support.DATA, split)
print("Is a PyTorch Dataset:", isinstance(train_dataset, Dataset))
print("Training examples:", len(train_dataset))
print("Validation examples:", len(validation_dataset))
print("Reserved test examples:", len(test_dataset))

original_index = int(split["train"][0])
raw_pixels = train_dataset.dataset.data[original_index]
image, label = train_dataset[0]

print("Stored pixel shape:", tuple(raw_pixels.shape))
print("Stored pixel dtype:", raw_pixels.dtype)
print("Stored pixel range:", raw_pixels.min().item(), raw_pixels.max().item())
print("Transformed image shape:", tuple(image.shape))
print("Transformed image dtype:", image.dtype)
print("Transformed pixel range:", image.min().item(), image.max().item())
print("Single class label:", label, "| Python type:", type(label).__name__)

def make_loaders(train_data, validation_data, batch_size=64, seed=0):
    # TODO: construct and return (training_loader, validation_loader).
    raise NotImplementedError("make_loaders")

check("data", globals())

train_loader, validation_loader = make_loaders(
    train_data=train_dataset,
    validation_data=validation_dataset,
    batch_size=64,
    seed=0,
)
images, labels = next(iter(train_loader))
print("Image batch shape:", tuple(images.shape))
print("Image dtype:", images.dtype)
print("Label batch shape:", tuple(labels.shape))
print("Label dtype:", labels.dtype)
results["data"] = {
    "train_size": len(train_dataset),
    "validation_size": len(validation_dataset),
    "batch_shape": tuple(images.shape),
    "image_dtype": str(images.dtype),
    "label_dtype": str(labels.dtype),
}

support.show_images(images=images, labels=labels)
plt.show()

# Cache images for the short diagnostic experiments only.
probe_train_data = support.cache_dataset(dataset=train_dataset)
probe_validation_data = support.cache_dataset(dataset=validation_dataset)
print("Diagnostic image shape:", tuple(probe_train_data.tensors[0].shape))

first_two = torch.stack([train_dataset[0][0], train_dataset[1][0]], dim=0)
print("stacked:", first_two.shape)
print("single-image batch:", image.unsqueeze(0).shape)
print("for a grayscale plot:", image.squeeze(0).shape)

mini = torch.arange(6).reshape(2, 3)
print("original:\n", mini)
print("reshape(3,2):\n", mini.reshape(3, 2))
print("reshape(2,-1) shape:", mini.reshape(2, -1).shape)
print("transpose:\n", mini.transpose(0, 1))
print("channel-last shape (for some display tools):", images.permute(0, 2, 3, 1).shape)

def flatten_images(images):
    # TODO: preserve the example axis; do not hard-code batch size 64.
    raise NotImplementedError("flatten_images")


def linear_scores(x, weight, bias):
    # x: (B,D), weight: (C,D), bias: (C,) -> scores: (B,C).
    raise NotImplementedError("linear_scores")


def predict_classes(logits):
    # TODO: return a long tensor of shape (B,).
    raise NotImplementedError("predict_classes")

flat = flatten_images(images)
torch.testing.assert_close(flat[0].reshape(1, 28, 28), images[0])
print("Flattened:", flat.shape, "| one-image batch:", flatten_images(images[:1]).shape)
weight = torch.zeros(10, 784)
bias = torch.linspace(-0.2, 0.2, 10)
scores = linear_scores(flat, weight, bias)
print("Scores:", scores.shape, "| predictions:", predict_classes(scores)[:8])

def cross_entropy_manual(logits, y):
    # TODO: stable mean loss and its (B,C) gradient, including the factor 1/B.
    raise NotImplementedError("cross_entropy_manual")

check("predictions", globals())

large_scores = torch.tensor(
    [[1000.0, 1001.0, 999.0], [900.0, 898.0, 901.0]],
    dtype=torch.float64,
)
example_labels = torch.tensor([1, 0], dtype=torch.long)
loss, score_gradient = cross_entropy_manual(large_scores, example_labels)
library_loss = F.cross_entropy(large_scores, example_labels)
print(f"Your stable cross-entropy: {loss.item():.6f}")
print(f"Library cross-entropy: {library_loss.item():.6f}")
print(f"Absolute loss difference: {abs(loss.item() - library_loss.item()):.3e}")
print("Gradient sum for each example:", score_gradient.sum(dim=1))

# Supplied preview only: autograd generates this linear classifier's gradients.
preview = support.linear_preview(
    train_data=probe_train_data,
    validation_data=probe_validation_data,
    steps=30,
    seed=0,
)
print("Training updates:", preview["steps"])
print(f"Validation accuracy after the preview: {preview['validation_accuracy']:.2%}")

def combine_linear_layers(w1, b1, w2, b2):
    # TODO: return the single equivalent (weight,bias).
    raise NotImplementedError("combine_linear_layers")

probe_params = support.manual_parameters(
    dims=(4, 5, 3),
    seed=7,
    dtype=torch.float64,
)
probe_x = torch.randn(6, 4, dtype=torch.float64)
combined_weight, combined_bias = combine_linear_layers(
    w1=probe_params["W1"],
    b1=probe_params["b1"],
    w2=probe_params["W2"],
    b2=probe_params["b2"],
)
hidden_without_relu = linear_scores(probe_x, probe_params["W1"], probe_params["b1"])
two_layer_scores = linear_scores(
    hidden_without_relu, probe_params["W2"], probe_params["b2"]
)
one_layer_scores = linear_scores(probe_x, combined_weight, combined_bias)
collapse_error = (two_layer_scores - one_layer_scores).abs().max().item()

relu_inputs = torch.tensor([-1.0, 0.0, 1.0])
relu_outputs = F.relu(relu_inputs)
midpoint_gap = abs((relu_outputs[0] + relu_outputs[2]) / 2 - relu_outputs[1]).item()
results["representations"] = {
    "collapse_max_error": collapse_error,
    "midpoint_gap": midpoint_gap,
}
print(f"Two-linear-layer collapse error: {collapse_error:.3e}")
print("Three collinear inputs:", relu_inputs)
print("Their ReLU outputs:", relu_outputs)
print(f"Midpoint-identity gap after ReLU: {midpoint_gap:.3f}")

def manual_mlp(params, x, y=None):
    # x has shape (batch_size, input_features).
    # Step 1: compute hidden preactivations, ReLU activations, and class logits.
    hidden_scores = None  # TODO: first linear layer
    hidden_activations = None  # TODO: ReLU
    logits = None  # TODO: second linear layer
    raise NotImplementedError(
        "Complete the forward computation, then remove this line."
    )

    if y is None:
        return logits

    # Step 2: obtain mean loss and its derivative with respect to the logits.
    loss, score_gradient = cross_entropy_manual(logits, y)

    # Step 3: send the derivative back through the output weights and ReLU.
    hidden_gradient = None  # TODO: gradient with respect to hidden activations
    hidden_score_gradient = None  # TODO: apply the ReLU derivative mask

    # Step 4: reduce over examples to obtain the four parameter gradients.
    grads = {
        "W1": None,  # TODO: shape matches params['W1']
        "b1": None,  # TODO: shape matches params['b1']
        "W2": None,  # TODO: shape matches params['W2']
        "b2": None,  # TODO: shape matches params['b2']
    }
    return {"logits": logits, "loss": loss, "grads": grads}

value = torch.tensor(3.0, requires_grad=True)
loss = value.square()
loss.backward()
print("After one backward:", value.grad.item())

loss = value.square()  # A fresh forward computation builds a new graph.
loss.backward()
print("After two backwards:", value.grad.item())

value.grad = None
value.square().backward()
print("After clearing and recomputing:", value.grad.item())

check("manual", globals())

small_params = support.manual_parameters(
    dims=(4, 5, 3),
    seed=2,
    dtype=torch.float64,
)
small_x = torch.randn(6, 4, dtype=torch.float64)
small_y = torch.tensor([0, 1, 2, 0, 2, 1])
results["verification"] = support.verify_manual(
    manual_mlp, small_params, small_x, small_y
)
for quantity, maximum_error in results["verification"].items():
    print(f"{quantity:>20}: maximum absolute error = {maximum_error:.3e}")

finite_difference = support.finite_difference_probe(
    forward=manual_mlp,
    params=small_params,
    x=small_x,
    y=small_y,
)
print(f"Analytic derivative: {finite_difference['analytic']:.8f}")
print(f"Finite-difference estimate: {finite_difference['numerical']:.8f}")
print(f"Absolute derivative difference: {finite_difference['absolute_error']:.3e}")

class TorchMLP(nn.Module):
    def __init__(self, dims=(784, 64, 10), dropout=0.0, batchnorm=False):
        super().__init__()
        self.linears = nn.ModuleList()
        self.norms = nn.ModuleList()

        for input_features, output_features in zip(dims[:-1], dims[1:]):
            # TODO: append a registered nn.Linear layer to self.linears.
            pass

        for hidden_features in dims[1:-1]:
            # TODO: append BatchNorm1d or Identity to self.norms.
            pass

        self.dropout = nn.Dropout(p=dropout)
        raise NotImplementedError("Finish registering layers, then remove this line.")

    def forward(self, images):
        x = None  # TODO: flatten images while preserving the batch axis
        for linear_layer, normalization in zip(self.linears[:-1], self.norms):
            # TODO: linear -> normalization -> ReLU -> dropout.
            pass

        raise NotImplementedError("Return raw logits from the final linear layer.")

check("model", globals())

bridge_params = support.manual_parameters(seed=3)
bridge_model = TorchMLP()
bridge_model.eval()
support.load_manual_weights(model=bridge_model, params=bridge_params)

with torch.no_grad():
    module_logits = bridge_model(images)
manual_logits = manual_mlp(bridge_params, flatten_images(images))
torch.testing.assert_close(module_logits, manual_logits)
parameter_count = sum(parameter.numel() for parameter in bridge_model.parameters())
print("Registered parameter count:", parameter_count)
print("Expected parameter count:", 50890)
maximum_logit_difference = (module_logits - manual_logits).abs().max().item()
print(f"Maximum module/manual logit difference: {maximum_logit_difference:.3e}")

sequential_example = nn.Sequential(
    nn.Flatten(start_dim=1),
    nn.Linear(in_features=784, out_features=64),
    nn.ReLU(),
    nn.Linear(in_features=64, out_features=10),
)

def momentum_step(w, grad, state=None, lr=0.01, beta=0.9):
    # TODO: persistent velocity; return new tensors without mutating the inputs.
    raise NotImplementedError("momentum_step")

def adamw_step(
    w,
    grad,
    state=None,
    lr=0.001,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
):
    beta1, beta2 = betas

    # Step 1: read persistent state, or initialize zero moments and counter zero.
    old_m = None  # TODO: tensor with the same shape/dtype/device as w
    old_v = None  # TODO: tensor with the same shape/dtype/device as w
    old_t = None  # TODO: integer number of previous updates

    # Step 2: update the moments using grad alone; keep weight decay separate.
    new_m = None  # TODO
    new_v = None  # TODO
    new_t = None  # TODO

    # Step 3: apply bias correction and construct the adaptive gradient update.
    corrected_m = None  # TODO
    corrected_v = None  # TODO
    new_weight = None  # TODO: decoupled decay plus the adaptive update

    raise NotImplementedError("Complete the three stages, then remove this line.")
    return new_weight, {"m": new_m, "v": new_v, "t": new_t}

manual_weight = torch.tensor(1.0, dtype=torch.float64)
manual_state = None
library_weight = nn.Parameter(manual_weight.clone())
library_optimizer = torch.optim.AdamW(
    [library_weight],
    lr=0.1,
    betas=(0.5, 0.5),
    eps=1e-8,
    weight_decay=0.2,
)

for step_number, gradient_value in enumerate([2.0, 0.0], start=1):
    gradient = torch.tensor(gradient_value, dtype=torch.float64)
    manual_weight, manual_state = adamw_step(
        w=manual_weight,
        grad=gradient,
        state=manual_state,
        lr=0.1,
        betas=(0.5, 0.5),
        eps=1e-8,
        weight_decay=0.2,
    )
    library_weight.grad = gradient.clone()
    library_optimizer.step()
    library_state = library_optimizer.state[library_weight]

    torch.testing.assert_close(manual_weight, library_weight.detach())
    torch.testing.assert_close(manual_state["m"], library_state["exp_avg"])
    torch.testing.assert_close(manual_state["v"], library_state["exp_avg_sq"])
    assert manual_state["t"] == step_number

    print(
        f"Step {step_number}: weight={manual_weight.item():.6f}, "
        f"m={manual_state['m'].item():.6f}, v={manual_state['v'].item():.6f}, "
        f"counter={manual_state['t']}"
    )
print("Both parameter updates and persistent moments match PyTorch.")

def train_epoch(model, loader, optimizer, label_smoothing=0.0):
    model.train()
    model_device = next(model.parameters()).device
    total_objective = 0.0
    correct = 0
    count = 0
    updates = 0

    for images, labels in loader:
        images = images.to(model_device)
        labels = labels.to(model_device)

        # Step 1: clear old gradients.
        optimizer.zero_grad(set_to_none=True)

        # Step 2: compute raw logits and the mean training objective.
        logits = None  # TODO: call the model
        loss = None  # TODO: cross-entropy with the requested label_smoothing
        raise NotImplementedError(
            "Complete this minibatch update, then remove this line."
        )

        # Step 3: backpropagate, then update registered parameters.
        # TODO: call backward and optimizer.step in the correct order.

        # Step 4: log detached values, weighted by the actual batch size.
        batch_count = labels.shape[0]
        # TODO: update total_objective, correct, count, and updates.

    return {
        "objective": total_objective / count,
        "accuracy": correct / count,
        "count": count,
        "updates": updates,
    }

def evaluate(model, loader):
    incoming_training_mode = model.training
    model.eval()
    model_device = next(model.parameters()).device
    total_loss = 0.0
    correct = 0
    count = 0
    all_logits = []
    all_targets = []

    try:
        with torch.no_grad():
            for images, labels in loader:
                images = images.to(model_device)
                labels = labels.to(model_device)
                logits = None  # TODO: call the model

                # TODO: accumulate hard-label loss, correct predictions and count.
                # Use a sum of example losses, not an unweighted sum of batch means.
                # TODO: append logits and labels on CPU to their corresponding lists.
                raise NotImplementedError("Complete evaluation for one minibatch.")
    finally:
        # Restore the caller's mode even if a check reports an error.
        model.train(incoming_training_mode)

    return {
        "loss": total_loss / count,
        "accuracy": correct / count,
        "count": count,
        "logits": None,  # TODO: concatenate all_logits along the batch axis
        "targets": None,  # TODO: concatenate all_targets along the batch axis
    }

check("training", globals())

def fit_model(model, train_loader, validation_loader, config, epochs=6):
    optimizer = None  # TODO: AdamW using this model's parameters and config.
    training_evaluation_loader = DataLoader(
        train_loader.dataset,
        batch_size=train_loader.batch_size,
        shuffle=False,
        num_workers=0,
        drop_last=False,
    )

    # Supplied records for comparing the runs later.
    initial_state = support.state_copy(model)
    initial_state_digest = support.tensor_digest(initial_state)
    order_digest = support.loader_order_digest(train_loader, epochs=epochs)
    parameter_count = sum(p.numel() for p in model.parameters())
    history = []
    updates = 0
    best_loss = math.inf
    best_epoch = None
    best_state = None

    for epoch in range(1, epochs + 1):
        trained = None  # TODO: train_epoch, including config['label_smoothing'].
        training = None  # TODO: evaluate on training_evaluation_loader.
        validation = None  # TODO: evaluate on validation_loader.
        raise NotImplementedError("Complete the epoch and checkpoint steps.")

        updates += trained["updates"]
        history.append(
            {
                "epoch": epoch,
                "training_objective": trained["objective"],
                "training": {k: training[k] for k in ("loss", "accuracy", "count")},
                "validation": {k: validation[k] for k in ("loss", "accuracy", "count")},
            }
        )
        # TODO: if validation loss improves, update best_loss, best_epoch,
        # and an independent copy of model.state_dict() in best_state.

    # TODO: restore best_state into the model.
    model.eval()
    training = evaluate(model, training_evaluation_loader)
    validation = evaluate(model, validation_loader)
    return {
        "config": copy.deepcopy(config),
        "history": history,
        "training": {k: training[k] for k in ("loss", "accuracy", "count")},
        "validation": {k: validation[k] for k in ("loss", "accuracy", "count")},
        "updates": updates,
        "parameter_count": parameter_count,
        "initial_state": initial_state,
        "initial_state_digest": initial_state_digest,
        "order_digest": order_digest,
        "best_epoch": best_epoch,
        "best_state": best_state,
        "model": model,
        "train_loader": train_loader,
        "validation_loader": validation_loader,
        "training_evaluation_loader": training_evaluation_loader,
    }

check("fit", globals())

SANITY_CHECK_LR = (
    0.003  # TODO: choose a workable learning rate for the 32-example check.
)

results["tiny"] = support.tiny_subset_fit(
    model_class=TorchMLP,
    train_epoch=train_epoch,
    evaluate=evaluate,
    train_data=probe_train_data,
    lr=SANITY_CHECK_LR,
    steps=150,
    seed=0,
)
print(f"Initial tiny-subset loss: {results['tiny']['initial_loss']:.6f}")
print(f"Final tiny-subset loss: {results['tiny']['final_loss']:.6f}")
print(f"Final tiny-subset accuracy: {results['tiny']['accuracy']:.2%}")
print("Learning rate:", results["tiny"]["lr"])
print("Training updates:", results["tiny"]["steps"])

plt.figure(figsize=(6, 3))
plt.plot(results["tiny"]["history"])
plt.xlabel("Update")
plt.ylabel("Tiny-subset training loss")
plt.title("Can this model fit 32 examples?")
plt.tight_layout()
plt.show()

OPTIMIZER_RATES = {
    "sgd": 0.1,
    "momentum": 0.03,
    "adamw": 0.003,
}  # TODO: choose and inspect.

results["optimizers"] = support.optimizer_comparison(
    forward=manual_mlp,
    momentum_step=momentum_step,
    adamw_step=adamw_step,
    train_data=probe_train_data,
    rates=OPTIMIZER_RATES,
    steps=80,
    seed=0,
)
print(
    f"{'Optimizer':<12} {'Learning rate':>14} "
    f"{'Initial loss':>14} {'Final loss':>12} {'Updates':>9}"
)
for optimizer_name, run in results["optimizers"].items():
    print(
        f"{optimizer_name:<12} {run['lr']:>14.4f} "
        f"{run['initial_loss']:>14.4f} {run['final_loss']:>12.4f} {run['steps']:>9d}"
    )

support.plot_optimizer_curves(runs=results["optimizers"])
plt.show()

def he_initialize(model, seed=0):
    # TODO: reproducible hidden/output scaling, zero biases, reset any BatchNorm state.
    raise NotImplementedError("he_initialize")

symmetry = support.symmetry_demo()
print(
    f"Incoming-weight gradient row difference: {symmetry['incoming_gradient_row_gap']:.3e}"
)
print(
    f"Outgoing-weight gradient column difference: {symmetry['outgoing_gradient_column_gap']:.3e}"
)

results["initialization"] = support.initialization_investigation(
    model_class=TorchMLP,
    initializer=he_initialize,
    train_epoch=train_epoch,
    evaluate=evaluate,
    train_data=probe_train_data,
    steps=40,
    seed=0,
)
print(
    f"{'Condition':<10} {'Layer':>6} {'Activation std':>16} "
    f"{'Inactive fraction':>19} {'Gradient norm':>16}"
)
for condition in ["before", "after"]:
    for layer_metrics in results["initialization"][condition]:
        print(
            f"{condition:<10} {layer_metrics['layer']:>6d} "
            f"{layer_metrics['activation_std']:>16.3e} "
            f"{layer_metrics['inactive_fraction']:>18.2%} "
            f"{layer_metrics['gradient_norm']:>16.3e}"
        )
original_initialization_loss = results["initialization"]["before_loss"]
print(
    f"Loss after 40 updates, original initialization: {original_initialization_loss:.4f}"
)
print(
    f"Loss after 40 updates, repaired initialization: {results['initialization']['after_loss']:.4f}"
)

def batchnorm_calculation(x, gamma, beta, eps=1e-5):
    feature_mean = None  # TODO: mean across examples, one value per feature
    feature_variance = None  # TODO: population variance across examples
    normalized = None  # TODO: center and divide by standard deviation
    raise NotImplementedError("Apply learned scale and shift, then return the result.")

bn_layer, batch_a, batch_b = support.batchnorm_fixture()
print("Learned parameters:", [name for name, _ in bn_layer.named_parameters()])
print("Running-state buffers:", [name for name, _ in bn_layer.named_buffers()])
manual_bn = batchnorm_calculation(
    x=batch_a,
    gamma=bn_layer.weight,
    beta=bn_layer.bias,
    eps=bn_layer.eps,
)
library_bn_layer = copy.deepcopy(bn_layer)
library_bn_layer.train()
with torch.no_grad():
    library_bn = library_bn_layer(batch_a)
normalization_error = (manual_bn - library_bn).detach().abs().max().item()
print("Your normalized output:\n", manual_bn.detach())
print(f"Maximum normalization error: {normalization_error:.3e}")
results["batchnorm"] = {"manual_error": normalization_error}

training_layer_a = copy.deepcopy(bn_layer)
training_layer_b = copy.deepcopy(bn_layer)
training_layer_a.train()
training_layer_b.train()

buffers_before_training = {
    name: value.detach().clone() for name, value in training_layer_a.named_buffers()
}
with torch.no_grad():
    training_output_a = training_layer_a(batch_a)
    training_output_b = training_layer_b(batch_b)

training_query_difference = (
    (training_output_a[0] - training_output_b[0]).abs().max().item()
)
training_buffer_changes = {
    name: not torch.equal(buffers_before_training[name], value)
    for name, value in training_layer_a.named_buffers()
}
print(f"Query-output difference in training mode: {training_query_difference:.6f}")
for name, changed in training_buffer_changes.items():
    print(f"{name} changed under no_grad: {changed}")
results["batchnorm"]["train_companion_gap"] = training_query_difference
results["batchnorm"]["running_mean_changed_under_no_grad"] = training_buffer_changes[
    "running_mean"
]

evaluation_layer = training_layer_a
evaluation_layer.eval()
buffers_before_evaluation = {
    name: value.detach().clone() for name, value in evaluation_layer.named_buffers()
}
with torch.no_grad():
    evaluation_output_a = evaluation_layer(batch_a)
    evaluation_output_b = evaluation_layer(batch_b)

evaluation_query_difference = (
    (evaluation_output_a[0] - evaluation_output_b[0]).abs().max().item()
)
evaluation_buffers_unchanged = all(
    torch.equal(buffers_before_evaluation[name], value)
    for name, value in evaluation_layer.named_buffers()
)
print(f"Query-output difference in evaluation mode: {evaluation_query_difference:.3e}")
print("All running-state buffers unchanged:", evaluation_buffers_unchanged)
results["batchnorm"]["eval_companion_gap"] = evaluation_query_difference
results["batchnorm"]["eval_buffers_unchanged"] = evaluation_buffers_unchanged

check("stability", globals())

def inverted_dropout(x, mask, p, training=True):
    # TODO: fixed-mask training calculation; identity values at evaluation.
    raise NotImplementedError("inverted_dropout")

activation = torch.tensor([[1.0, -2.0, 3.0], [0.0, 4.0, -1.0]])
fixed_mask = torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
results["dropout"] = {
    "training": inverted_dropout(activation, fixed_mask, p=0.25, training=True),
    "evaluation": inverted_dropout(activation, fixed_mask, p=0.25, training=False),
}
print("Original activations:\n", activation)
print("Fixed keep-mask:\n", fixed_mask)
print("Training output:\n", results["dropout"]["training"])
print("Evaluation output:\n", results["dropout"]["evaluation"])

REGULARIZER = "dropout"  # TODO: choose dropout, weight_decay, or label_smoothing.
REGULARIZER_STRENGTH = 0.2  # TODO: choose a positive strength below 1.

def run_study(configs, train_data, validation_data, seed=0):
    runs = {}
    train_loaders = {}
    validation_loaders = {}

    for name, config in configs.items():
        torch.manual_seed(seed)
        model = None  # TODO: TorchMLP using dims, dropout and batchnorm from config.
        # TODO: initialize this model with he_initialize and the supplied seed.
        train_loader, validation_loader = None, None
        # TODO: make_loaders using config['batch_size'] and seed + 101.
        runs[name] = None  # TODO: fit_model for six epochs.
        raise NotImplementedError("Assemble each run, then remove this line.")
        train_loaders[name] = train_loader
        validation_loaders[name] = validation_loader

    selected = None  # TODO: name of the run with the lowest validation loss.

    # Supplied records for plots and comparison checks.
    validation_x, validation_y = support.loader_tensors(validation_loaders[selected])
    regularizer = next(
        k
        for k in ("dropout", "weight_decay", "label_smoothing")
        if configs["regularized"][k] > 0
    )
    return {
        "runs": runs,
        "selected": selected,
        "models": {name: run["model"] for name, run in runs.items()},
        "validation_x": validation_x,
        "validation_y": validation_y,
        "regularizer": regularizer,
        "strength": configs["regularized"][regularizer],
        "train_loaders": train_loaders,
        "validation_loaders": validation_loaders,
        "train_data": train_data,
        "validation_data": validation_data,
    }

baseline_config = {
    "dims": (784, 64, 10),
    "dropout": 0.0,
    "batchnorm": False,
    "optimizer": "adamw",
    "lr": 0.003,
    "weight_decay": 0.0,
    "label_smoothing": 0.0,
    "batch_size": 64,
}
configs = {
    "baseline": baseline_config,
    "regularized": {**baseline_config, REGULARIZER: REGULARIZER_STRENGTH},
    "batchnorm": {**baseline_config, "batchnorm": True},
}
results["study"] = run_study(configs, train_dataset, validation_dataset, seed=0)

print(
    f"{'Run':<14} {'Epoch':>6} {'Train loss':>12} {'Val loss':>12} "
    f"{'Val accuracy':>14} {'Updates':>9}"
)
for run_name, run in results["study"]["runs"].items():
    print(
        f"{run_name:<14} {run['best_epoch']:>6d} "
        f"{run['training']['loss']:>12.4f} {run['validation']['loss']:>12.4f} "
        f"{run['validation']['accuracy']:>13.2%} {run['updates']:>9d}"
    )
print("Selected by validation loss:", results["study"]["selected"])

support.plot_study(study=results["study"])
plt.show()

def finish_study(study, test_loader, path="a1_selected_model.pt"):
    selected_run = study["runs"][study["selected"]]
    checkpoint = {
        "config": selected_run["config"],
        "state_dict": selected_run["best_state"],
    }
    # TODO: save checkpoint to path, then load it into saved.
    saved = None
    config = None  # TODO: read the saved configuration.
    loaded = None  # TODO: create a fresh TorchMLP with that configuration.
    # TODO: load the saved state into loaded.
    raise NotImplementedError("Complete save/reload, prediction and test evaluation.")
    loaded.eval()

    # Supplied comparison: before/after reload predictions should match.
    test_x, test_y = support.loader_tensors(test_loader)
    source_model = study["models"][study["selected"]]
    source_model.eval()
    with torch.no_grad():
        source_logits = source_model(test_x[:16])
        loaded_logits = loaded(test_x[:16])
        single_image_logits = None  # TODO: predict test_x[0] with a batch axis.

    test_metrics = None  # TODO: evaluate the reloaded model on test_loader.
    predictions = None  # TODO: predicted class for each row of test_metrics['logits'].
    error_mask = None  # TODO: compare predictions with test_metrics['targets'].

    # Supplied records for reporting and plotting.
    study["test"] = {k: test_metrics[k] for k in ("loss", "accuracy", "count")}
    return {
        "max_logit_error": (source_logits - loaded_logits).abs().max().item(),
        "single_image_shape": tuple(single_image_logits.shape),
        "loaded_model": loaded,
        "source_model": source_model,
        "source_logits": source_logits,
        "loaded_logits": loaded_logits,
        "config": config,
        "path": str(path),
        "error_mask": error_mask,
        "error_images": test_x[error_mask][:8],
        "error_targets": test_y[error_mask][:8],
        "error_predictions": predictions[error_mask][:8],
        "test_x": test_x,
        "test_y": test_y,
        "test_metrics": test_metrics,
        "test_loader": test_loader,
    }

check("generalization", globals())

test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)
results["checkpoint"] = finish_study(results["study"], test_loader)
print(
    f"Reloaded-model logit difference: {results['checkpoint']['max_logit_error']:.3e}"
)
print("Single-image output shape:", results["checkpoint"]["single_image_shape"])
print(f"Selected-model test loss: {results['study']['test']['loss']:.4f}")
print(f"Selected-model test accuracy: {results['study']['test']['accuracy']:.2%}")
print("Checkpoint file:", results["checkpoint"]["path"])

support.show_images(
    images=results["checkpoint"]["error_images"],
    labels=results["checkpoint"]["error_targets"],
    predictions=results["checkpoint"]["error_predictions"],
)
plt.show()

Part	Main work	Points
1. Images to predictions	Data loading, shapes, scores and stable loss	12
2. Hidden layers and gradients	ReLU, manual backpropagation, autograd and modules	20
3. Optimization and training	Momentum, AdamW, training and evaluation	22
4. Initialization and normalization	Initialization, diagnosis and BatchNorm training	14
5. Generalization	Dropout, controlled comparisons and model selection	12
Problem set	Four calculation/short-answer questions in a separate PDF	20

Run	Setting
Baseline	784 → 64 → 10; AdamW learning rate 0.003; no dropout, decay or smoothing
Regularized	Choose one of dropout, weight decay or label smoothing and its strength
BatchNorm	Add BatchNorm before the hidden ReLU; explicit regularizers remain zero

Assignment 1 · Neural network fundamentals¶

0. Set up your PyTorch workspace¶

Tensor basics¶

What is a tensor?¶

Practice: indexing¶

Practice: broadcasting¶

1. From images to predictions · 12 points¶

1.1 Load images and build batches · 4 points¶

1.2 Give each image ten class scores · 4 points¶

1.3 Turn scores into a stable learning signal · 4 points¶

2. Hidden layers and backpropagation · 20 points¶

2.1 What does a hidden layer add? · 2 points¶

2.2 Implement the hidden layer and chain rule · 12 points¶

2.3 Check against autograd¶

2.4 Express the network as a module · 6 points¶

3. Optimization and training · 22 points¶

3.1 Make gradients change the weights · 3 points¶

3.2 Implement AdamW · 4 points¶

3.3 Write a reusable training loop · 5 points¶

3.4 Evaluate a fixed model · 4 points¶

3.5 Train across epochs and keep the best model · 2 points¶

3.6 Check trainability on a tiny set · 2 points¶

3.7 Compare update rules · 2 points¶

4. Initialization and normalization · 14 points¶

4.1 Apply a fan-in rule · 4 points¶

4.2 Diagnose and repair initialization · 4 points¶

4.3 Calculate BatchNorm and inspect its modes · 3 points¶

4.4 Train a BatchNorm MLP · 3 points¶

5. Generalization and model selection · 12 points¶

5.1 Compute dropout · 3 points¶

5.2 Compare three candidates using validation · 6 points¶

5.3 Save, reload and test your model · 3 points¶

Before submitting¶