The list transformer made each forward operation and backward derivative visible. Replacing those loops with a tensor library changes how the computation is expressed. It should preserve the computation you intended. Here you port the existing transformer block to PyTorch and compare both its outputs and derivatives before considering a training loop.
When you need this. You have a small model whose equations and manual gradients you understand. You want tensor operations and automatic differentiation, but a correct output shape would not reveal a reversed matrix, missing causal mask or differently averaged loss. A numerical bridge checks those choices against the earlier implementation.
A tensor is an array with an explicit shape, numeric type and device. Autograd records tensor operations involving parameters marked requires_grad=True; calling loss.backward() computes their derivatives. This lab keeps the model on CPU in float64 so numerical comparisons are easier to interpret. It compares one synthetic next-token sequence, rather than training a language model.
The saved run compares logits, mean cross-entropy and all 18 parameter-gradient tensors, containing 268 elements. Both tensor paths pass atol=1e-10 and rtol=1e-8. These tolerances allow floating-point rounding; they are not a promise of identical results across devices or library releases.
1. Carry the existing parameters into tensors
Download the complete lab and open its pytorch-transformer-bridge/ directory. The lab keeps an unchanged copy of the earlier model.py under vendor/. It creates parameters with the same seeded list initializer, then converts each to a separate float64 CPU leaf tensor. There is no pretrained-model download.
python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txtUse the PyTorch installation guide if the pinned package is unavailable for your platform. The recorded library version is retained in the raw JSON. Choosing another installation requires rerunning the checks; this lab has not compared other devices or dtypes.
The input IDs are [1,3,1,5,2], with next-token targets [3,1,5,2,6]. These are authored synthetic IDs in a vocabulary of nine. Token 1 repeats, so its embedding gradient must accumulate contributions from multiple positions. The position table has eight rows, leaving three unused rows whose gradients should be zero.
The model has width four and an MLP hidden width of seven. A projection matrix stays shaped (input_width, output_width), matching the list implementation’s x @ W. PyTorch’s tensor multiplication accepts that orientation directly; introducing nn.Linear would require accounting for its different stored weight orientation.
2. Match the attention and normalization contracts
The tensor input has shape (5, 4). Each row receives a token embedding plus its learned position embedding. Layer normalization precedes attention and the MLP, using the existing gains, biases and eps=1e-5. Residual additions remain around both sublayers. The MLP uses ReLU, and the output projection produces logits shaped (5, 9).
The explicit path computes softmax(Q @ K.T / sqrt(4)) @ V, replacing future-position scores with negative infinity before softmax. The SDPA path adds batch and head dimensions to make (1, 1, 5, 4), calls scaled_dot_product_attention, then removes those two singleton dimensions. It passes is_causal=True, the same scale and dropout_p=0.0, and selects the math backend.
The SDPA API defines the scale, causal mask and dropout arguments. Its dropout behavior follows the argument even during evaluation. Setting a module to evaluation mode alone would not disable dropout inside this function. This lab does not test a fused GPU kernel.
Read the complete tensor bridge
"""Compare the unchanged list transformer with CPU float64 PyTorch autograd."""
import argparse
import hashlib
import json
import math
import platform
import sys
from pathlib import Path
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from vendor import model
HERE = Path(__file__).resolve().parent
ATOL, RTOL = 1e-10, 1e-8
CONFIG = {'vocab': 9, 'd': 4, 'hidden': 7, 'max_len': 8, 'seed': 31}
IDS, TARGETS = [1, 3, 1, 5, 2], [3, 1, 5, 2, 6]
def params():
return model.init_params(**CONFIG)
def tensors(p):
return {name: torch.tensor(value, dtype=torch.float64, device='cpu', requires_grad=True)
for name, value in p.items()}
def inputs(p, ids, targets):
vocab = len(p['tok_emb'])
if (not ids or len(ids) > len(p['pos_emb']) or len(targets) != len(ids)
or any(type(i) is not int or not 0 <= i < vocab for i in [*ids, *targets])):
raise ValueError('require nonempty in-range integer IDs, matching targets and available positions')
def attention(q, k, v, path='explicit', *, causal=True, scale=None):
scale = 1 / math.sqrt(q.shape[-1]) if scale is None else scale
if path == 'explicit':
scores = (q @ k.T) * scale
if causal:
allowed = torch.ones(q.shape[0], k.shape[0], dtype=torch.bool).tril()
scores = scores.masked_fill(~allowed, float('-inf'))
return scores.softmax(dim=-1) @ v
if path != 'sdpa':
raise ValueError('unknown attention path')
# Shape: (batch=1, heads=1, positions, width). Force the documented math backend.
with sdpa_kernel(backends=[SDPBackend.MATH]):
mixed = F.scaled_dot_product_attention(q[None, None], k[None, None], v[None, None],
dropout_p=0.0, is_causal=causal, scale=scale)
return mixed[0, 0]
def forward(p, ids, targets, path='explicit', *, causal=True, scale=None, eps=model.EPS):
inputs(p, ids, targets)
x0 = p['tok_emb'][ids] + p['pos_emb'][:len(ids)]
norm = lambda x, name: F.layer_norm(x, (CONFIG['d'],), p[name+'_g'], p[name+'_b'], eps=eps)
h1 = norm(x0, 'ln1')
q, k, v = (h1 @ p[name] for name in ('Wq', 'Wk', 'Wv'))
x1 = x0 + attention(q, k, v, path, causal=causal, scale=scale) @ p['Wo']
h2 = norm(x1, 'ln2')
x2 = x1 + F.relu(h2 @ p['W1'] + p['b1']) @ p['W2'] + p['b2']
logits = norm(x2, 'lnf') @ p['W_out'] + p['b_out']
loss = F.cross_entropy(logits, torch.tensor(targets, dtype=torch.long), reduction='mean')
return logits, loss
def difference(actual, expected):
a = torch.as_tensor(actual, dtype=torch.float64).detach()
b = torch.as_tensor(expected, dtype=torch.float64).detach()
if a.shape != b.shape or not torch.isfinite(a).all() or not torch.isfinite(b).all():
raise ValueError('comparison needs equal finite shapes')
return {'max_abs_error': (a-b).abs().max().item(),
'passed': bool(torch.allclose(a, b, atol=ATOL, rtol=RTOL))}
def compare(path, p=None):
p = params() if p is None else p
manual_logits, cache = model.forward(p, IDS)
manual_loss, dlogits = model.loss_and_dlogits(manual_logits, TARGETS)
manual_grads = model.backward(p, cache, dlogits)
tensor_p = tensors(p)
logits, loss = forward(tensor_p, IDS, TARGETS, path)
loss.backward()
gradients = {name: difference(tensor_p[name].grad, manual_grads[name]) for name in p}
return {'logits': difference(logits, manual_logits), 'loss': difference(loss, manual_loss),
'loss_value': loss.item(), 'parameter_gradients': gradients,
'compared_gradient_elements': sum(t.numel() for t in tensor_p.values()),
'nonzero_manual_gradient_tensors': sum(any(v != 0 for v in flatten(manual_grads[n])) for n in p),
'passed': all(x['passed'] for x in gradients.values()) and
difference(logits, manual_logits)['passed'] and difference(loss, manual_loss)['passed']}
def flatten(value):
if isinstance(value, list):
return [n for child in value for n in flatten(child)]
return [value]
def failures():
p = tensors(params())
reference, _ = forward(p, IDS, TARGETS)
rows = {}
for name, kwargs in [('no_causal_mask', {'causal':False}), ('unscaled_attention', {'scale':1.0}),
('wrong_layernorm_epsilon', {'eps':0.1})]:
changed, _ = forward(p, IDS, TARGETS, **kwargs)
rows[name] = difference(changed, reference)
refused = []
for ids, targets in [([], []), ([1, 2], [2]), ([1, 9], [2, 3]), ([True], [1])]:
try:
forward(p, ids, targets)
except ValueError as error:
refused.append({'ids':ids, 'targets':targets, 'error':str(error)})
else:
raise AssertionError('invalid input accepted')
return {'mismatch_probes':rows, 'refused_inputs':refused}
def record():
runs = {path: compare(path) for path in ('explicit', 'sdpa')}
probes = failures()
assert all(run['passed'] for run in runs.values())
assert all(not row['passed'] for row in probes['mismatch_probes'].values())
files = ['bridge.py', 'vendor/model.py', 'test_bridge.py', 'requirements.txt']
return {'execution_mode':'CPU float64 forward and full parameter-backward comparison; no training update',
'command':['python','bridge.py','--record'], 'exit_status':0,
'config':CONFIG, 'ids':IDS, 'targets':TARGETS, 'dtype':'float64', 'device':'cpu',
'dropout_p':0.0, 'is_causal':True, 'attention_scale':1/math.sqrt(CONFIG['d']),
'layernorm_eps':model.EPS, 'loss_reduction':'mean', 'sdpa_backend':'MATH',
'tolerance':{'atol':ATOL,'rtol':RTOL}, 'comparisons':runs, 'failures':probes,
'environment':{'python':platform.python_version(),'torch':torch.__version__,'platform':platform.platform()},
'sources':{name:hashlib.sha256((HERE/name).read_bytes()).hexdigest() for name in files}}
def main():
parser=argparse.ArgumentParser();parser.add_argument('--record',action='store_true');args=parser.parse_args()
torch.set_num_threads(1)
result=record()
print('CPU float64; one head; five positions; dropout=0; causal scale=1/sqrt(4)')
for path, row in result['comparisons'].items():
print(f"{path}: logits max error={row['logits']['max_abs_error']:.3e}; loss error={row['loss']['max_abs_error']:.3e}; "
f"gradient max error={max(x['max_abs_error'] for x in row['parameter_gradients'].values()):.3e}; "
f"{len(row['parameter_gradients'])} tensors/{row['compared_gradient_elements']} elements; passed={row['passed']}")
for name, probe in result['failures']['mismatch_probes'].items():
print(f"{name}: logits max error={probe['max_abs_error']:.3e}; within tolerance={probe['passed']}")
print('Invalid inputs refused:',len(result['failures']['refused_inputs']))
if args.record:
(HERE/'result.json').write_text(json.dumps(result,indent=2,allow_nan=False)+'\n')
if __name__=='__main__':main()A missing causal mask. no_causal_mask lets earlier positions read later tokens. The saved probe changes logits by 1.142e-3, outside tolerance. Restore causal masking before comparing gradients. test_no_causal_mask_changes_outputs checks this failure.
A missing attention scale. unscaled_attention uses scale 1 instead of 1/sqrt(4). The saved difference is 2.205e-6. The operation runs, but it computes a different attention distribution. test_missing_scale_changes_outputs requires the mismatch to remain visible.
3. Compare outputs and every parameter gradient
Run the comparison, then the tests:
python bridge.py
python -m unittest -v test_bridge.pyThe loss consumes unnormalized logits and integer targets. Both implementations average cross-entropy across the five positions. The tensor paths call backward() on fresh parameter tensors, so a previous call cannot silently accumulate another run’s gradients. Every tensor is compared element by element with the manual backward pass, not just a selected weight.
CPU float64; one head; five positions; dropout=0; causal scale=1/sqrt(4)
explicit: logits max error=1.388e-17; loss error=4.441e-16; gradient max error=8.327e-17; 18 tensors/268 elements; passed=True
sdpa: logits max error=1.388e-17; loss error=4.441e-16; gradient max error=8.327e-17; 18 tensors/268 elements; passed=True
no_causal_mask: logits max error=1.142e-03; within tolerance=False
unscaled_attention: logits max error=2.205e-06; within tolerance=False
wrong_layernorm_epsilon: logits max error=1.042e-01; within tolerance=False
Invalid inputs refused: 4The acceptance rule is abs(actual − reference) <= atol + rtol * abs(reference). An absolute term handles values close to zero; a relative term scales with larger values. The maximum absolute gradient error in the saved run is below 1e-16 for each tensor path. The raw result includes each tensor’s measured error, input IDs, settings, source hashes, command, exit status and runtime. Run python report.py in the extracted lab to generate an HTML report from those saved results; this does not run the model again.
A different normalization epsilon. wrong_layernorm_epsilon changes epsilon to 0.1. Its measured logit difference is 0.1042. test_wrong_layernorm_epsilon_changes_outputs detects that change. A tensor rewrite must preserve normalization details as well as attention.
Malformed input. Empty IDs, mismatched target lengths, an out-of-range ID and a Boolean ID are refused before model execution. test_invalid_ids_and_target_lengths_refused checks all four cases. A correctly shaped output would not compensate for comparing different next-token targets.
Read the executable checks
import hashlib
import json
import unittest
import torch
import bridge as b
class BridgeTests(unittest.TestCase):
def test_saved_run_still_matches_implementation_sources(self):
record=json.loads((b.HERE/'result.json').read_text())
for name, expected in record['sources'].items():
self.assertEqual(hashlib.sha256((b.HERE/name).read_bytes()).hexdigest(),expected)
self.assertTrue(all(row['passed'] for row in record['comparisons'].values()))
def test_explicit_outputs_loss_and_all_gradients(self):
r=b.compare('explicit');self.assertTrue(r['passed'])
self.assertEqual(len(r['parameter_gradients']),len(b.params()))
self.assertEqual(r['nonzero_manual_gradient_tensors'],len(b.params()))
def test_sdpa_outputs_loss_and_all_gradients(self):
self.assertTrue(b.compare('sdpa')['passed'])
def test_repeated_token_embedding_accumulates_gradient(self):
p=b.tensors(b.params());_,loss=b.forward(p,b.IDS,b.TARGETS);loss.backward()
self.assertGreater(p['tok_emb'].grad[1].abs().max().item(),0)
self.assertTrue(b.compare('sdpa')['parameter_gradients']['tok_emb']['passed'])
def test_unused_embedding_rows_and_positions_have_zero_gradient(self):
p=b.tensors(b.params());_,loss=b.forward(p,b.IDS,b.TARGETS);loss.backward()
self.assertEqual(p['tok_emb'].grad[8].abs().max().item(),0)
self.assertEqual(p['pos_emb'].grad[len(b.IDS):].abs().max().item(),0)
def test_no_causal_mask_changes_outputs(self):
self.assertFalse(b.failures()['mismatch_probes']['no_causal_mask']['passed'])
def test_missing_scale_changes_outputs(self):
self.assertFalse(b.failures()['mismatch_probes']['unscaled_attention']['passed'])
def test_wrong_layernorm_epsilon_changes_outputs(self):
self.assertFalse(b.failures()['mismatch_probes']['wrong_layernorm_epsilon']['passed'])
def test_invalid_ids_and_target_lengths_refused(self):
self.assertEqual(len(b.failures()['refused_inputs']),4)
def test_causal_future_token_cannot_change_earlier_logits(self):
for path in ('explicit','sdpa'):
p=b.tensors(b.params());a,_=b.forward(p,b.IDS,b.TARGETS,path)
ids=b.IDS[:-1]+[7];c,_=b.forward(p,ids,b.TARGETS,path)
self.assertTrue(torch.equal(a[:-1],c[:-1]));self.assertFalse(torch.equal(a[-1],c[-1]))
def test_bad_comparison_shapes_and_nonfinite_values_refused(self):
for a,c in [([1],[1,2]),([float('nan')],[1])]:
with self.assertRaises(ValueError):b.difference(a,c)
def test_unknown_attention_path_refused(self):
with self.assertRaises(ValueError):b.forward(b.tensors(b.params()),b.IDS,b.TARGETS,'unknown')
if __name__=='__main__':unittest.main()
How the comparison earns its conclusion
All 18 parameter tensors have nonzero manual gradients somewhere in this fixture. The comparison also checks unused embedding and position rows, which should stay zero. A repeated token exercises gradient accumulation into one embedding row. Changing the final token must leave earlier logits unchanged in each causal tensor path. These checks probe distinct parts of the port rather than relying on its output shape.
The handwritten implementation and PyTorch implement the same intended equations through different backward mechanisms. Agreement on this small fixture supports the port at this boundary. It does not evaluate generated text, optimizer updates, long contexts, low precision or accelerator kernels. Keep this CPU comparison as a regression test before expanding those parts.
Check your understanding
Read the explanation
First align the equations and settings, then compare outputs and gradients under an explicit numerical tolerance.
What’s next
Return to the training pipeline to follow data windows, parameter updates and training/validation separation. This bridge performs no parameter update; adding one requires its own optimizer-state and training checks.
References
Official PyTorch 2.14 documentation checked October 24, 2026. The linked sections specify the API behavior used here; the saved local run supplies the numerical comparisons.
- SDPA: causal masking, scale, dropout and backend selection.
- Layer normalization: normalized dimensions and epsilon.
- Cross-entropy: logits, class-index targets and mean reduction.