LLM Developer DocsLLM Docs

Loading the search index ...

↑ ↓ move, Enter open, Esc close

GitHub

ops.py

Differentiable operations with hand-written backward functions

Where it fits

11 of 12 functions, classes and methods here are reachable from train.main, and 3 from talk.main.

Pipeline stageStepExplained on
BlocksLayerNorm 1layernorm
BlocksLayerNorm 2layernorm
OutputFinal LayerNormfull model
Loss and updateSoftmax and cross-entropysoftmax cross entropy
Loss and updateBackwardhow the model learns

Imports and used by

Imports from this repository
None
Other imports
numpy
Imported by
generate.py, tests/test_ops.py, train.py, transformer.py
Used by, in the pipeline
generate.py, train.py, transformer.py
Used by, in tests
tests/test_ops.py

Mathematics used

ConceptWhat this file doesLinesStatus
Xavier uniform initialisationComputes it4, 5Implemented
Matrix multiplicationComputes it8Derived from implementation
Linear layerComputes it8Implemented
Linear layer, backwardComputes it11, 12, 13Implemented
ReLUComputes it17Implemented
ReLU, backwardComputes it20Implemented
Softmax per rowComputes it23, 24Implemented
SoftmaxComputes it23, 24Implemented
TemperatureCalls it24Implemented
Softmax, backwardComputes it27Implemented
Attention backwardComputes it27Implemented
Cross-entropyComputes it30, 31, 33, 34Implemented
Cross-entropy, backwardComputes it38, 39, 40Implemented
Deriving the backward formulaComputes it43, 44, 45, 46, 51, 52, 53, 55, 56Derived from implementation
LayerNormComputes it43, 44, 45, 46, 47Implemented
LayerNorm, backwardComputes it51, 52, 53, 55, 56Implemented
PerplexityComputes it60Implemented

Functions and constants

Source

ops.py60 linesGitHub
import numpy as npdef xavier_uniform(fan_in, fan_out, rng):    a = np.sqrt(6.0 / (fan_in + fan_out))    return rng.uniform(-a, a, size=(fan_in, fan_out))def linear_forward(X, W, b):    return X @ W + bdef linear_backward(X, W, dY):    dW = X.T @ dY    db = dY.sum(axis=0)    dX = dY @ W.T    return dX, dW, dbdef relu_forward(X):    return np.maximum(0.0, X)def relu_backward(Z, dA):    return dA * (Z > 0)def softmax_forward(x):    e = np.exp(x - np.max(x, axis=-1, keepdims=True))    return e / e.sum(axis=-1, keepdims=True)def softmax_backward(y, dY):    return y * (dY - np.sum(y * dY, axis=-1, keepdims=True))def cross_entropy_forward(logits, targets):    z = logits - logits.max(axis=-1, keepdims=True)    logp = z - np.log(np.exp(z).sum(axis=-1, keepdims=True))    T = len(targets)    loss = -logp[np.arange(T), targets].mean()    return float(loss), np.exp(logp) def cross_entropy_backward(probs, targets):    T = len(targets)    dZ = probs.copy()    dZ[np.arange(T), targets] -= 1.0    return dZ / Tdef layer_norm_forward(X, gamma, beta, epsilon = 1e-5):    mu = X.mean(axis=-1, keepdims=True)    var = ((X - mu) ** 2).mean(axis=-1, keepdims=True)    s = 1.0 / np.sqrt(var + epsilon)    Xhat = (X - mu) * s    return gamma * Xhat + beta, (Xhat, s)def layer_norm_backward(dY, cache, gamma):    Xhat, s = cache    dXhat = dY * gamma    dX = s * (dXhat - dXhat.mean(axis=-1, keepdims=True)                - Xhat * (dXhat * Xhat).mean(axis=-1, keepdims=True))        dgamma = (dY * Xhat).sum(axis=0)    dbeta = dY.sum(axis=0)    return dX, dgamma, dbetadef perplexity(loss):    return float(np.exp(loss))