Activations and normalization

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
%matplotlib inline

# read all words
words = open("names.txt", "r").read().splitlines()
print(words[:8])
print(len(words))

# Build the vocabulary of characters and mappings to/from integers
chars = sorted(list(set(''.join(words))))
stoi = {s : i+1 for i, s in enumerate(chars)}
stoi['.'] = 0
itos = {i:s for s, i in stoi.items()}
vocab_size = len(itos)
print(itos)
print(vocab_size)

# build the dataset
# context length: how many characters do we take to predict the next one?
block_size = 3

def build_dataset(words):
    X, Y = [], []

    for w in words:
        context = [0] * block_size
        for ch in w + '.':
            ix = stoi[ch]
            X.append(context)
            Y.append(ix)
            context = context[1:] + [ix] #crop on left and append

    X = torch.tensor(X)
    Y = torch.tensor(Y)
    print(X.shape, Y.shape)
    return X, Y

import random
random.seed(42)
random.shuffle(words)
n1 = int(0.8 * len(words))
n2 = int(0.9 * len(words))

Xtr, Ytr = build_dataset(words[:n1]) #80%
Xdev, Ydev = build_dataset(words[n1:n2]) #10%
Xte, Yte = build_dataset(words[n2:]) #10%

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g)
b1 = torch.randn(n_hidden, generator = g)
W2 = torch.randn((n_hidden, vocab_size), generator = g)
b2 = torch.randn(vocab_size, generator = g)

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

1/27.0
torch.tensor(1/27.0).log()
-torch.tensor(1/27.0).log()

# let's see a smaller 4-dimensional example of this issue
logits = torch.tensor([0.0, 0.0, 0.0, 0.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
probs, loss

# lets experiment now and see how the loss is affected when values are diff w.r.t one another
logits = torch.tensor([0.0, 0.0, 5.0, 0.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
probs, loss

# in this case, the initialization of correct probability to the correct label is by chance, thus we are getting low loss
# look at index 2: 5 and we are setting loss for index 2 as well, now look at another example

# this case is much more likely than randomly guessing correct weights, the loss is much higher
logits = torch.tensor([0.0, 5.0, 0.0, 0.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
probs, loss

# logits = torch.tensor([0.0, 5.0, 0.0, 0.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
logits = torch.randn(4)
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
logits, probs, loss

# logits = torch.tensor([0.0, 5.0, 0.0, 0.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
logits = torch.randn(4) * 100
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
logits, probs, loss

# SETTING ALL TO 1: OK CUZ OF NORMALIZATION INSIDE SOFTMAX LATER

logits = torch.tensor([1.0, 1.0, 1.0, 1.0]) #suppose we have only 4 characters, and the logits that come out of neural net are very close to 0
probs = torch.softmax(logits, dim=0) #when we take a softmax of all 0's, we get probabilities that are diffused distribution
loss = -probs[2].log() #here, it doesn't matter if label is 0,1,2,3 cuz its a uniform distribution, they all sum to 1
logits, probs, loss

# after optimization, our loss looks something like this:
plt.plot(lossi)

# this decorator disables gradient tracking
# when this decorator is used on top of a function
# torch assumes that whatever happens in this function never requires any gradients
# so it wont keep track of anything to keep track of gradients in anticipation of backward pass
# has the same effect as setting required_grad of all tensors here to false
# makes more efficient, cuz we let torch know that we wont call .backward() on any computation
# so torch doesnt have to maintain a graph under the hood
# you can also use a context manager with torch.no_grad

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    h = torch.tanh(embcat @ W1 + b1) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

# sample from the model
g = torch.Generator().manual_seed(2147483647 + 10)

for _ in range(20):

    out = []
    context = [0] * block_size #initialize with all ...

    while True:

        # forward pass the neural net
        emb = C[torch.tensor((context))] #(1, block_size, n_embd)
        h = torch.tanh(emb.view(1, -1) @ W1 + b1)
        logits = h @ W2 + b2
        probs = F.softmax(logits, dim = 1)

        # sample from the distribution
        ix = torch.multinomial(probs, num_samples = 1, generator=g).item()

        # shift the context window and track the samples
        context = context[1:] + [ix]
        out.append(ix)

        # if we sample the special '.' then break
        if ix == 0:
            break

    print(''.join(itos[i] for i in out)) #decode and print the generated word

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #break after the very first iteration itself

# the loss is way too high, lets examine the logits randomly
# okay, so the logits take on extreme values, this creates fake confidence in incorrect answers thus increasing the loss
logits[5]
# so: these logits should be much much closer to zero

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g)
b1 = torch.randn(n_hidden, generator = g)
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #break after the very first iteration itself

# now, the logits are much closer to 0
logits[5]

# let's now run the optimization with this new initialization, and see what losses we record
# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

plt.plot(lossi)

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    h = torch.tanh(embcat @ W1 + b1) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g)
b1 = torch.randn(n_hidden, generator = g)
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #break after the very first iteration itself

# LETS visualize the tensor h, notice that many of the elements are 1 or -1
print(h)
# recall that tanh is a squashing function, it takes arbitrary numbers and squashes them between the range of 1 and -1 smoothly
# lets look at the histogram of h to get a better idea of the distribution of the values inside this tensor
print(h.shape) #32 examples(our batch size) with 200 activations for each of the 32 examples

# we can stretch h out into a one large vector
print(h.view(-1).shape)
# then convert it into one large python list of floats
# h.view(-1).tolist()

plt.hist(h.view(-1).tolist(), 50);
# semicolon to supress some output we don't want

# from the histogram, we see that most of the values take on -1 and +1
# so our tanh is very very active, and to see why so, we can look at the preactivations
# hpreact = embcat @ W1 + b1 #hidden layer pre-activation
# h = torch.tanh(hpreact) #hidden layer
# logits = h @ W2 + b2 #output layer

# distribution of the preactivations is very broad: they take nums between -15 and 15, so in torch.tanh everything is being capped to be in range -1 to 1
# and lots of numbers take extreme value(-1, 1)
plt.hist(hpreact.view(-1).tolist(), 50);

plt.figure(figsize=(20,10))
plt.imshow(h.abs() > 0.99, cmap = 'gray', interpolation='nearest')
# we want to take absolute value of h and see how often it is in the flat line of tanh graph
# h.abs() > 0.99 is a boolean tensor: white if true and black if false
# graph shows 32 examples and 200 neurons: majority is white, they're in flat tail and the backward gradient would get destroyed
# we want to see more black instead of white

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * 0.1 #added this
b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #break after the very first iteration itself

plt.hist(h.view(-1).tolist(), 50);
# semicolon to supress some output we don't want
# wow, look at the difference in output: it was more towards 1 and -1 previously

plt.hist(hpreact.view(-1).tolist(), 50);

plt.figure(figsize=(20,10))
plt.imshow(h.abs() > 0.99, cmap = 'gray', interpolation='nearest')

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * 0.2 #changed this from 0.1 to 0.2
b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 1000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #break after the very first iteration itself

plt.hist(h.view(-1).tolist(), 50);
# semicolon to supress some output we don't want
# now, this is a nice distribution

plt.hist(hpreact.view(-1).tolist(), 50);

plt.figure(figsize=(20,10))
plt.imshow(h.abs() > 0.99, cmap = 'gray', interpolation='nearest')

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    h = torch.tanh(embcat @ W1 + b1) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

x = torch.randn(1000, 10) # 1000 x 10 tensor: 1000 examples that are 10-dimensional each
w = torch.randn(10, 200) # 10 x 200 tensor #weight layer also initialized using Gaussian/Normal/Bell shaped distribution, #200 neurons in the hidden layer look at 10 input each
y = x @ w #1000 x 200 tensor #get the preactivations of the neurons #for now, forget bias and non-linearity
print(x.mean(), x.std())
print(y.mean(), y.std())
plt.figure(figsize=(20, 5))
plt.subplot(121)
plt.hist(x.view(-1).tolist(), 50, density=True);
plt.subplot(122)
plt.hist(y.view(-1).tolist(), 50, density=True);

# initially mean/s.d. of x and w = 0/1. we are curious about how the nature of y will be 
# so the mean and s.d. of the gaussian/normal distribution 
# s.d: just the spread of the gaussian/normal distribution
# we notice that y.mean stays the same: about 0 cuz its a symmetric operation, but the s.d. has expanded to 3 from 1
# but we don't want that, we want the neural net to have mostly similar activations, roughly unit Gaussian throughout the neural net
# so, now the qn is: how do we scale/multiply the W with to preserve the distribution to remain gaussian, i.e. to ensure that s.d. of y will also be 1 later after matrix multiplication

x = torch.randn(1000, 10) # 1000 x 10 tensor: 1000 examples that are 10-dimensional each
w = torch.randn(10, 200) / (10 ** 0.5) #200 neurons in the hidden layer look at 10 input each
y = x @ w #1000 x 200 tensor #get the preactivations of the neurons #for now, forget bias and non-linearity
print(x.mean(), x.std())
print(y.mean(), y.std())
plt.figure(figsize=(20, 5))
plt.subplot(121)
plt.hist(x.view(-1).tolist(), 50, density=True);
plt.subplot(122)
plt.hist(y.view(-1).tolist(), 50, density=True);

# create 10k numbers and look at the std
# its the spread of the Gaussian/Normal distribution with mean 0 and s.d. 1
torch.randn(1000).std()

# multiply it by 0.2 to get s.d of 0.2
(torch.randn(1000) * 0.2).std()

# what we're trying to show is the num we multiply with there ends up being the s.d. of the gaussian

# gain / (fan_in)**(-n)
# gain divided by square root of fan_in
# the fan_in for W1 was n_embd * block_size which is 30. yes n_embd=10 but we had 3 of them and concatenated them so we should have used 30 instead of 10 in mini example of x, w and y = x @ w
(5/3) / (30 ** 0.5)
# this number is what we want our standard deviation to be

(5/3) / ((n_embd * block_size) ** 0.5)

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * (5/3) / ((n_embd * block_size) ** 0.5)
b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself

parameters = [C, W1, b1, W2, b2]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

plt.plot(lossi)

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    h = torch.tanh(embcat @ W1 + b1) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

hpreact.shape

# calculate mean across the 0th dimension (vertically down) and keepdims = True to make it easily broadcastable
hpreact.mean(0, keepdim=True)
hpreact.std(0, keepdim=True)

hpreact.mean(0, keepdim=True).shape

hpreact.std(0, keepdim=True).shape

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * (5/3) / ((n_embd * block_size) ** 0.5)
b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself
bngain = torch.ones((1, n_hidden))
bnbias = torch.zeros((1, n_hidden))

parameters = [C, W1, b1, W2, b2, bngain, bnbias]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation

    hpreact = bngain * (hpreact - hpreact.mean(0, keepdim=True)) / hpreact.std(0, keepdim=True) + bnbias
    
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    hpreact = embcat @ W1 + b1
    
    hpreact = bngain * (hpreact - hpreact.mean(0, keepdim=True)) / hpreact.std(0, keepdim=True) + bnbias

    h = torch.tanh(hpreact) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

# calibrate/adjust the batch norm at the end of training:

with torch.no_grad():
    # pass the training set through
    emb = C[Xtr]
    embcat = emb.view(emb.shape[0], -1)
    hpreact = embcat @ W1 + b1

    # measure the mean/std over the entire training set
    bmean = hpreact.mean(0, keepdim=True)
    bnstd = hpreact.std(0, keepdim=True)

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    hpreact = embcat @ W1 + b1
    
    hpreact = bngain * (hpreact - bmean) / bnstd + bnbias

    h = torch.tanh(hpreact) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * (5/3) / ((n_embd * block_size) ** 0.5)
b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself
bngain = torch.ones((1, n_hidden))
bnbias = torch.zeros((1, n_hidden))
bnmean_running = torch.zeros((1, n_hidden))
bnstd_running = torch.ones((1, n_hidden))

parameters = [C, W1, b1, W2, b2, bngain, bnbias]
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors
    hpreact = embcat @ W1 + b1 #hidden layer pre-activation #will be roughly unit Gaussian in the first iteration

    bnmeani = hpreact.mean(0, keepdim=True) #batch normalization mean for the ith iteration #will be roughly 0 for the first iteration
    bnstdi = hpreact.std(0, keepdim=True) #will be roughly 1 for the first iteration
    hpreact = bngain * (hpreact - bnmeani) / bnstdi  + bnbias 

    # we don't want pytorch to build a graph for update of mean/sd
    # this update is separate than the gradient based optimization during backprop
    with torch.no_grad():
        # it will be mostly what it already is, but receive a small update based on the direction of what the current mean/sd is
        bnmean_running = 0.999 * bnmean_running + 0.001 * bnmeani
        bnstd_running = 0.999 * bnstd_running + 0.001 * bnstdi
    
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

# calibrate/adjust the batch norm at the end of training:

with torch.no_grad():
    # pass the training set through
    emb = C[Xtr]
    embcat = emb.view(emb.shape[0], -1)
    hpreact = embcat @ W1 + b1

    # measure the mean/std over the entire training set
    bnmean = hpreact.mean(0, keepdim=True)
    bnstd = hpreact.std(0, keepdim=True)

bnmean

bnmean_running

bnstd

bnstd_running

@torch.no_grad()
def split_loss(split):

    x, y = {
        'train': (Xtr, Ytr),
        'val': (Xdev, Ydev),
        'test': (Xte, Yte),
    }[split] #depending on train, val or test as input, index and get the correct split

    # forward pass of network: evaluate the loss and print it
    emb = C[x] # N x block_size x n_embd
    embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size x n_embd) to allow matrix mul
    hpreact = embcat @ W1 + b1
    
    hpreact = bngain * (hpreact - bnmean_running) / bnstd_running + bnbias

    h = torch.tanh(hpreact) # N x n_hidden
    logits = h @ W2 + b2 # N x vocab_size
    loss = F.cross_entropy(logits, y)
    print(split, loss.item())

split_loss('train')
split_loss('val')

# MLP revisited
n_embd = 10 #dimensionality of the character embedding vectors
n_hidden = 200 #number of neurons in hidden layer of MLP

g = torch.Generator().manual_seed(2147483647) #use same seed for reproducibility
C = torch.randn((vocab_size, n_embd), generator = g)
W1 = torch.randn((n_embd * block_size, n_hidden), generator = g) * (5/3) / ((n_embd * block_size) ** 0.5)
# b1 = torch.randn(n_hidden, generator = g) * 0.01 #added this
W2 = torch.randn((n_hidden, vocab_size), generator = g) * 0.01 #scaling it down so it is near to 0 and logits will not be very extreme thus reducing loss at initialization itself
b2 = torch.randn(vocab_size, generator = g) * 0 #not adding anything at initialization to get low loss during first iteration itself
bngain = torch.ones((1, n_hidden))
bnbias = torch.zeros((1, n_hidden))

parameters = [C, W1, W2, b2, bngain, bnbias] #removed b1 cuz it was stopping training cuz it was None
print(sum(p.nelement() for p in parameters))
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    embcat = emb.view(emb.shape[0], -1) #concatenate the vectors

    # Linear layer
    hpreact = embcat @ W1 # + b1 #hidden layer pre-activation #will be roughly unit Gaussian in the first iteration #removed b1 from here and above parameters list too

    # Batch normalization (batchnorm) layer
    # ----------------------------------------------------------------------------------------------------------------------------------
    bnmeani = hpreact.mean(0, keepdim=True) #batch normalization mean for the ith iteration #will be roughly 0 for the first iteration
    bnstdi = hpreact.std(0, keepdim=True) #will be roughly 1 for the first iteration
    hpreact = bngain * (hpreact - bnmeani) / bnstdi  + bnbias 

    # we don't want pytorch to build a graph for update of mean/sd
    # this update is separate than the gradient based optimization during backprop
    with torch.no_grad():
        # it will be mostly what it already is, but receive a small update based on the direction of what the current mean/sd is
        bnmean_running = 0.999 * bnmean_running + 0.001 * bnmeani
        bnstd_running = 0.999 * bnstd_running + 0.001 * bnstdi
    # ----------------------------------------------------------------------------------------------------------------------------------

    # non-linearity
    h = torch.tanh(hpreact) #hidden layer
    logits = h @ W2 + b2 #output layer
    loss = F.cross_entropy(logits, Yb) #loss function

    # backward pass
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())

    break #after first iteration itself

# Let's train a deeper network

class Linear:

    def __init__(self, fan_in, fan_out, bias=True):
        self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5
        self.bias = torch.zeros(fan_out) if bias else None

    def __call__(self, x):
        self.out = x @ self.weight
        if self.bias is not None:
            self.out += self.bias
        return self.out

    def parameters(self):
        return [self.weight] + ([] if self.bias is None else [self.bias])

class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum = 0.1):
        self.eps = eps
        self.momentum = momentum
        self.training = True

        # parameters(trained with backprop)
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)

        # buffers (trained with a running 'momentum update')
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)

    def __call__(self, x):

        # calculate the forward pass
        if self.training:
            xmean = x.mean(0, keepdim=True) #batch mean
            xvar = x.var(0, keepdim=True, unbiased=True) #batch variance
        else:
            xmean = self.running_mean
            xvar = self.running_var

        xhat = (x - xmean) / torch.sqrt(xvar + self.eps) #normalize to unit variance
        self.out = self.gamma * xhat + self.beta

        # update the buffers
        if self.training:
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
        return self.out

    def parameters(self):
        return [self.gamma, self.beta]

class Tanh:
    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out

    def parameters(self):
        return []

n_embd = 10 #dimensionality of character embedding vectors
n_hidden = 100 #number of neurons in hidden layer of MLP
g = torch.Generator().manual_seed(2147483647) #for reproducibility

C = torch.randn((vocab_size, n_embd), generator=g)
layers = [
    Linear(n_embd * block_size, n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), Tanh(),
    Linear(n_hidden, vocab_size),
]

with torch.no_grad():

    # last layer: make less confident
    layers[-1].weight *= 0.1

    # all other layers: apply gain
    for layer in layers[:-1]:
        if isinstance(layer, Linear):
            layer.weight *= 5/3 #1

parameters = [C] + [p for layer in layers for p in layer.parameters()]
print(sum(p.nelement() for p in parameters)) # no of parameters in total
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []
ud = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    x = emb.view(emb.shape[0], -1) #concatenate the vectors

    for layer in layers:
        x = layer(x)
    loss = F.cross_entropy(x, Yb) #loss function

    # backward pass
    for layer in layers:
        layer.out.retain_grad() #after debug: would take out retain grad
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())
    with torch.no_grad():
        ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters])
    
    if i >= 1000: #for the last 2 graphs, need to go through more iterations
        break #after first iteration itself, obviously remove this after debug

# activation statistics in forward pass(if gain set to 1 instead of 5/3)

# visualize histograms of forward pass activations specifically at tanh layers
plt.figure(figsize=(20,4)) #width and height of the plot
legends = []

for i, layer in enumerate(layers[:-1]): #while iterating over all layer,s exclude only the output softmax layer
    if isinstance(layer, Tanh): #tanh cuz they have finite output -1 to +1, easy to visualize
        t = layer.out
        # we define saturation as > 0.97, which is tail of tanh, and it stops gradients
        print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach(), hy.detach())
        legends.append(f'layer {i} ( {layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# activation statistics in forward pass(if gain set to 5/3)

# visualize histograms of forward pass activations specifically at tanh layers
plt.figure(figsize=(20,4)) #width and height of the plot
legends = []

for i, layer in enumerate(layers[:-1]): #while iterating over all layer,s exclude only the output softmax layer
    if isinstance(layer, Tanh): #tanh cuz they have finite output -1 to +1, easy to visualize
        t = layer.out
        # we define saturation as > 0.97, which is tail of tanh, and it stops gradients
        print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach(), hy.detach())
        legends.append(f'layer {i} ( {layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# activation statistics in backward pass

# visualize histograms
plt.figure(figsize=(20,4)) #width and height of the plot
legends = []

for i, layer in enumerate(layers[:-1]): #exclude the output layer
    if isinstance(layer, Tanh):
        t = layer.out.grad
        print('layer %d (%10s): mean %+f, std %e' % (i, layer.__class__.__name__, t.mean(), t.std()))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach(), hy.detach())
        legends.append(f'layer {i} {layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# visualize histograms
plt.figure(figsize=(20, 4)) # width and height of the plot
legends = []
for i,p in enumerate(parameters):
  t = p.grad
  if p.ndim == 2:
    print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
    hy, hx = torch.histogram(t, density=True)
    plt.plot(hx[:-1].detach(), hy.detach())
    legends.append(f'{i} {tuple(p.shape)}')
plt.legend(legends)
plt.title('weights gradient distribution');

plt.figure(figsize=(20, 4))
legends = []
for i,p in enumerate(parameters):
  if p.ndim == 2:
    plt.plot([ud[j][i] for j in range(len(ud))])
    legends.append('param %d' % i)
plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plot
plt.legend(legends);

# Let's train a deeper network

class Linear:

    def __init__(self, fan_in, fan_out, bias=True):
        self.weight = torch.randn((fan_in, fan_out), generator=g) / fan_in**0.5
        self.bias = torch.zeros(fan_out) if bias else None

    def __call__(self, x):
        self.out = x @ self.weight
        if self.bias is not None:
            self.out += self.bias
        return self.out

    def parameters(self):
        return [self.weight] + ([] if self.bias is None else [self.bias])

class BatchNorm1d:

    def __init__(self, dim, eps=1e-5, momentum = 0.1):
        self.eps = eps
        self.momentum = momentum
        self.training = True

        # parameters(trained with backprop)
        self.gamma = torch.ones(dim)
        self.beta = torch.zeros(dim)

        # buffers (trained with a running 'momentum update')
        self.running_mean = torch.zeros(dim)
        self.running_var = torch.ones(dim)

    def __call__(self, x):

        # calculate the forward pass
        if self.training:
            xmean = x.mean(0, keepdim=True) #batch mean
            xvar = x.var(0, keepdim=True, unbiased=True) #batch variance
        else:
            xmean = self.running_mean
            xvar = self.running_var

        xhat = (x - xmean) / torch.sqrt(xvar + self.eps) #normalize to unit variance
        self.out = self.gamma * xhat + self.beta

        # update the buffers
        if self.training:
            with torch.no_grad():
                self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * xmean
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * xvar
        return self.out

    def parameters(self):
        return [self.gamma, self.beta]

class Tanh:
    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out

    def parameters(self):
        return []

n_embd = 10 #dimensionality of character embedding vectors
n_hidden = 100 #number of neurons in hidden layer of MLP
g = torch.Generator().manual_seed(2147483647) #for reproducibility

C = torch.randn((vocab_size, n_embd), generator=g)
layers = [
    Linear(n_embd * block_size, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden, n_hidden), BatchNorm1d(n_hidden), Tanh(),
    Linear(n_hidden, vocab_size), BatchNorm1d(vocab_size),
]

with torch.no_grad():

    # last layer: make less confident
    # layers[-1].weight *= 0.1
    # since the last layer is batchnorm, we won't change the weight to make the softmax less confident
    # we'll change gamma: variable that multiplicatively interacts with output of normalization

    layers[-1].gamma *= 0.1

    # all other layers: apply gain
    for layer in layers[:-1]:
        if isinstance(layer, Linear):
            layer.weight *= 5/3 #1

parameters = [C] + [p for layer in layers for p in layer.parameters()]
print(sum(p.nelement() for p in parameters)) # no of parameters in total
for p in parameters:
    p.requires_grad = True

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []
ud = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    x = emb.view(emb.shape[0], -1) #concatenate the vectors

    for layer in layers:
        x = layer(x)
    loss = F.cross_entropy(x, Yb) #loss function

    # backward pass
    for layer in layers:
        layer.out.retain_grad() #after debug: would take out retain grad
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())
    with torch.no_grad():
        ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters])
    
    if i >= 1000: #for the last 2 graphs, need to go through more iterations
        break #after first iteration itself, obviously remove this after debug

# this looks really good, cuz of normalization of batch before every tanh
# std is roughly 0.65, saturation 2%, and roughly equal s.d. throughout the entire layers
# lol, so we had forgotten to put batch normalization for first layer, and it was so weird and different from others

# activation statistics in forward pass(if gain set to 5/3)

# visualize histograms of forward pass activations specifically at tanh layers
plt.figure(figsize=(20,4)) #width and height of the plot
legends = []

for i, layer in enumerate(layers[:-1]): #while iterating over all layer,s exclude only the output softmax layer
    if isinstance(layer, Tanh): #tanh cuz they have finite output -1 to +1, easy to visualize
        t = layer.out
        # we define saturation as > 0.97, which is tail of tanh, and it stops gradients
        print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach(), hy.detach())
        legends.append(f'layer {i} ( {layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# activation statistics in backward pass

# visualize histograms
plt.figure(figsize=(20,4)) #width and height of the plot
legends = []

for i, layer in enumerate(layers[:-1]): #exclude the output layer
    if isinstance(layer, Tanh):
        t = layer.out.grad
        print('layer %d (%10s): mean %+f, std %e' % (i, layer.__class__.__name__, t.mean(), t.std()))
        hy, hx = torch.histogram(t, density=True)
        plt.plot(hx[:-1].detach(), hy.detach())
        legends.append(f'layer {i} {layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# visualize histograms
plt.figure(figsize=(20, 4)) # width and height of the plot
legends = []
for i,p in enumerate(parameters):
  t = p.grad
  if p.ndim == 2:
    print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
    hy, hx = torch.histogram(t, density=True)
    plt.plot(hx[:-1].detach(), hy.detach())
    legends.append(f'{i} {tuple(p.shape)}')
plt.legend(legends)
plt.title('weights gradient distribution');

plt.figure(figsize=(20, 4))
legends = []
for i,p in enumerate(parameters):
  if p.ndim == 2:
    plt.plot([ud[j][i] for j in range(len(ud))])
    legends.append('param %d' % i)
plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plot
plt.legend(legends);

# same optimization as previously done
max_steps = 200000
batch_size = 32
lossi = []
ud = []

for i in range(max_steps):

    # minibatch construct
    ix = torch.randint(0, Xtr.shape[0], (batch_size, ), generator = g)
    Xb, Yb = Xtr[ix], Ytr[ix] #batch X, Y

    # forward pass
    emb = C[Xb] #embed the characters into vectors
    x = emb.view(emb.shape[0], -1) #concatenate the vectors

    for layer in layers:
        x = layer(x)
    loss = F.cross_entropy(x, Yb) #loss function

    # backward pass
    for layer in layers:
        layer.out.retain_grad() #after debug: would take out retain grad
    for p in parameters:
        p.grad = None
    loss.backward()

    # update
    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay
    for p in parameters:
        p.data += -lr * p.grad

    # track stats
    if i % 10000 == 0: #print every once in a while
        print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')
    lossi.append(loss.log10().item())
    with torch.no_grad():
        ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters])

# visualize histograms
plt.figure(figsize=(20, 4)) # width and height of the plot
legends = []
for i, layer in enumerate(layers[:-1]): # note: exclude the output layer
  if isinstance(layer, Tanh):
    t = layer.out
    print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100))
    hy, hx = torch.histogram(t, density=True)
    plt.plot(hx[:-1].detach(), hy.detach())
    legends.append(f'layer {i} ({layer.__class__.__name__}')
plt.legend(legends);
plt.title('activation distribution')

# visualize histograms
plt.figure(figsize=(20, 4)) # width and height of the plot
legends = []
for i, layer in enumerate(layers[:-1]): # note: exclude the output layer
  if isinstance(layer, Tanh):
    t = layer.out.grad
    print('layer %d (%10s): mean %+f, std %e' % (i, layer.__class__.__name__, t.mean(), t.std()))
    hy, hx = torch.histogram(t, density=True)
    plt.plot(hx[:-1].detach(), hy.detach())
    legends.append(f'layer {i} ({layer.__class__.__name__}')
plt.legend(legends);
plt.title('gradient distribution')

# visualize histograms
plt.figure(figsize=(20, 4)) # width and height of the plot
legends = []
for i,p in enumerate(parameters):
  t = p.grad
  if p.ndim == 2:
    print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
    hy, hx = torch.histogram(t, density=True)
    plt.plot(hx[:-1].detach(), hy.detach())
    legends.append(f'{i} {tuple(p.shape)}')
plt.legend(legends)
plt.title('weights gradient distribution');

plt.figure(figsize=(20, 4))
legends = []
for i,p in enumerate(parameters):
  if p.ndim == 2:
    plt.plot([ud[j][i] for j in range(len(ud))])
    legends.append('param %d' % i)
plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plot
plt.legend(legends);

@torch.no_grad() # this decorator disables gradient tracking
def split_loss(split):
  x,y = {
    'train': (Xtr, Ytr),
    'val': (Xdev, Ydev),
    'test': (Xte, Yte),
  }[split]
  emb = C[x] # (N, block_size, n_embd)
  x = emb.view(emb.shape[0], -1) # concat into (N, block_size * n_embd)
  for layer in layers:
    x = layer(x)
  loss = F.cross_entropy(x, y)
  print(split, loss.item())

# put layers into eval mode
for layer in layers:
  layer.training = False
split_loss('train')
split_loss('val')

emb.shape

# sample from the model
g = torch.Generator().manual_seed(2147483647 + 10)

for _ in range(20):

    out = []
    context = [0] * block_size #initialize with all ...

    while True:

        # forward pass the neural net
        # emb = C[torch.tensor((context))] #(1, block_size, n_embd)
        # x = emb.view(emb.shape[0], -1) #concatenate the vectors
        # x = emb.view(x.shape[0], -1)

        emb = C[torch.tensor([context])]   # (1, block_size, n_embd)
        x = emb.view(1, -1)                # (1, block_size*n_embd)
        
        for layer in layers:
            x = layer(x)
        logits = x
        probs = F.softmax(logits, dim=1)

        # sample from the distribution
        ix = torch.multinomial(probs, num_samples=1, generator=g).item()

        # shift the context window and track the samples
        context = context[1:] + [ix]
        out.append(ix)
        # if we sample the special '.' token, break
        if ix == 0:
            break
        
    print(''.join(itos[i] for i in out)) # decode and print the generated word

# BatchNorm forward pass as a widget

from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets
import scipy.stats as stats
import numpy as np

def normshow(x0):
  
  g = torch.Generator().manual_seed(2147483647+1)
  x = torch.randn(5, generator=g) * 5
  x[0] = x0 # override the 0th example with the slider
  mu = x.mean()
  sig = x.std()
  y = (x - mu)/sig

  plt.figure(figsize=(10, 5))
  # plot 0
  plt.plot([-6,6], [0,0], 'k')
  # plot the mean and std
  xx = np.linspace(-6, 6, 100)
  plt.plot(xx, stats.norm.pdf(xx, mu, sig), 'b')
  xx = np.linspace(-6, 6, 100)
  plt.plot(xx, stats.norm.pdf(xx, 0, 1), 'r')
  # plot little lines connecting input and output
  for i in range(len(x)):
    plt.plot([x[i],y[i]], [1, 0], 'k', alpha=0.2)
  # plot the input and output values
  plt.scatter(x.data, torch.ones_like(x).data, c='b', s=100)
  plt.scatter(y.data, torch.zeros_like(y).data, c='r', s=100)
  plt.xlim(-6, 6)
  # title
  plt.title('input mu %.2f std %.2f' % (mu, sig))

interact(normshow, x0=(-30,30,0.5));

# Linear: activation statistics of forward and backward pass

g = torch.Generator().manual_seed(2147483647)

a = torch.randn((1000,1), requires_grad=True, generator=g)          # a.grad = b.T @ c.grad
b = torch.randn((1000,1000), requires_grad=True, generator=g)       # b.grad = c.grad @ a.T
c = b @ a
loss = torch.randn(1000, generator=g) @ c
a.retain_grad()
b.retain_grad()
c.retain_grad()
loss.backward()
print('a std:', a.std().item())
print('b std:', b.std().item())
print('c std:', c.std().item())
print('-----')
print('c grad std:', c.grad.std().item())
print('a grad std:', a.grad.std().item())
print('b grad std:', b.grad.std().item())

# Linear + BatchNorm: activation statistics of forward and backward pass

g = torch.Generator().manual_seed(2147483647)

n = 1000
# linear layer ---
inp = torch.randn(n, requires_grad=True, generator=g)
w = torch.randn((n, n), requires_grad=True, generator=g) # / n**0.5
x = w @ inp
# bn layer ---
xmean = x.mean()
xvar = x.var()
out = (x - xmean) / torch.sqrt(xvar + 1e-5)
# ----
loss = out @ torch.randn(n, generator=g)
inp.retain_grad()
x.retain_grad()
w.retain_grad()
out.retain_grad()
loss.backward()

print('inp std: ', inp.std().item())
print('w std: ', w.std().item())
print('x std: ', x.std().item())
print('out std: ', out.std().item())
print('------')
print('out grad std: ', out.grad.std().item())
print('x grad std: ', x.grad.std().item())
print('w grad std: ', w.grad.std().item())
print('inp grad std: ', inp.grad.std().item())