Tiny Full Forward-Pass Fixtures for Minimal Models

Full report

On this page

Executive summary

A tiny forward-pass fixture is most valuable when it is small enough to run instantly, strict enough to fail loudly, and deterministic enough to serve as a parity gate. In practice, the best progression is: start with a single-layer MLP for the first smoke test, add a tiny CNN if you need tensor-layout and convolution coverage, and then add a single Transformer block if your real target is an LLM-style execution path. PyTorch and TensorFlow are the easiest places to author reference fixtures quickly, JAX is excellent when you want explicit control over randomness and a compiled math path, and ONNX Runtime is the right inference-oriented fixture when you want to validate deployment/runtime behavior or bridge into a C# or ONNX-first stack. PyTorch explicitly warns that full reproducibility is not guaranteed across releases, commits, platforms, or CPU-vs-GPU runs even with identical seeds, so the safest parity baseline remains fixed versions + fixed hardware + CPU + float32 + deterministic flags.

For a fixture whose only job is “prove one whole forward pass is wired correctly,” use float32 activations, dropout = 0, inference/eval mode, and an exact input shape chosen in advance. For timing, always separate first-call overhead from steady-state single-pass timing by doing one warm-up run before the timed run. That matters especially for JAX because jax.jit introduces compilation behavior, and it matters for ONNX Runtime because you create and configure an InferenceSession before executing run().

If your larger goal is a managed-runtime or C# inference engine that eventually replaces a heavier local runtime, a tiny full forward-pass fixture is exactly the right milestone: it gives you a narrow numerical contract before you expand into tokenization, KV cache, quantization, or scheduler work. That aligns directly with the uploaded roadmap context.

Fixture design

A robust fixture should do six things in one pass: construct the model, seed all randomness deterministically, create a deterministic input tensor, run exactly one forward pass, record the output tensor and elapsed time, and validate a few invariant properties such as shape, finiteness, and non-degenerate values. PyTorch provides torch.manual_seed() and deterministic-algorithm controls; TensorFlow provides tf.keras.utils.set_random_seed() and deterministic op execution; JAX uses explicit PRNG keys for deterministic random sequences; ONNX Runtime executes a prebuilt graph through InferenceSession.run() and lets you control providers and threading at the session level.

flowchart TD
    A[Select tiny model architecture] --> B[Fix exact shape and dtype]
    B --> C[Seed framework and input RNG]
    C --> D[Build model in inference mode]
    D --> E[Create deterministic input tensor]
    E --> F[Warm-up run]
    F --> G[Timed forward pass]
    G --> H[Capture output tensor]
    H --> I[Validate shape, dtype, finiteness, non-degeneracy]
    I --> J[Store golden outputs and timing metadata]

The simplest command-line workflow is to keep one environment per framework, pin versions, and run one small script per fixture. Current official install guides show PyTorch’s local installer flow, TensorFlow’s pip installation with tensorflow or tensorflow[and-cuda] depending platform, JAX’s CPU/GPU/TPU extras, and ONNX Runtime’s CPU or GPU packages. ONNX Runtime also documents the C# package command dotnet add package Microsoft.ML.OnnxRuntime, which is useful if you later port the same fixture logic into a .NET validation harness.

python -m venv .venv
source .venv/bin/activate   # Windows: .\.venv\Scripts\activate

pip install torch
pip install tensorflow
pip install -U jax
pip install onnx onnxruntime numpy

python fixture_pytorch.py
python fixture_tensorflow.py
python fixture_jax.py
python fixture_onnxruntime.py

# If you later mirror the ONNX Runtime fixture in .NET:
dotnet add package Microsoft.ML.OnnxRuntime

Minimal architectures, shapes, and dtypes

For a first fixture, the single-layer MLP is the best baseline because it exercises parameter loading, affine math, output capture, and shape validation while avoiding layout differences and attention-specific complexity. Both PyTorch Linear and TensorFlow Dense apply a dense transform on the last dimension, so a compact input like [2, 4] with dtype float32 is ideal; a corresponding tiny output such as [2, 3] gives enough room for sanity checks without introducing noise from larger matrices.

A tiny CNN is the next-best fixture when you want to validate convolution semantics and tensor layout. PyTorch’s Conv2d uses channel-first shapes like (N, C, H, W), while TensorFlow Conv2D defaults to channel-last (batch, height, width, channels). For a genuinely tiny but meaningful input, use an 8×8 grayscale image in float32: [1, 1, 8, 8] for PyTorch-style layouts and [1, 8, 8, 1] for TensorFlow-style layouts. A practical full tiny-CNN fixture is Conv2D/Conv2d(1→2, kernel=3) -> ReLU -> global average pool -> Dense(3), which yields a compact final output of [1, 3].

A single Transformer block is the right minimal analogue when your real target is a decoder-style model stack. PyTorch’s TransformerEncoderLayer with batch_first=True expects input/output tensors shaped (batch, seq, feature), and TensorFlow MultiHeadAttention documents query/value tensors shaped (B, T, dim). A very small but still representative configuration is batch = 1, sequence length = 4, model width = 8, heads = 2, feed-forward width = 16, dtype = float32, dropout = 0.0, giving an input shape of [1, 4, 8] and output shape [1, 4, 8]. This is the best “second-stage” fixture after an MLP if you eventually care about LLM-layer parity.

The exact numbers above are recommendations rather than framework requirements. They are chosen to stay small enough for instant execution, large enough to catch obvious shape and math errors, and simple enough to hand-check if needed. For cross-framework parity, float32 is the safest default because it is universally supported and avoids the extra complexity of mixed precision or accelerator-specific reduced-precision paths. PyTorch and TensorFlow layer docs both describe the relevant input-shape conventions, and JAX and ONNX Runtime both work cleanly with explicitly created float32 arrays for this purpose.

Runnable examples

The four examples below are deliberately small, runnable, and bounded to a single forward pass. Each one performs deterministic setup, a warm-up run, a timed run, output capture, and basic assertions.

PyTorch example using a tiny Transformer block

PyTorch documents torch.manual_seed() for seeding all devices and torch.use_deterministic_algorithms(True) for deterministic algorithm selection when available. TransformerEncoderLayer with batch_first=True gives the convenient (batch, seq, feature) shape used in the fixture below.

import time, torch
seed = 1234
torch.manual_seed(seed)
torch.use_deterministic_algorithms(True)
torch.set_num_threads(1)

m = torch.nn.TransformerEncoderLayer(
    d_model=8, nhead=2, dim_feedforward=16, dropout=0.0,
    batch_first=True, dtype=torch.float32).eval()

x = torch.randn(1, 4, 8, dtype=torch.float32)

with torch.inference_mode():
    _ = m(x)  # warm-up
    t0 = time.perf_counter()
    y = m(x)
    ms = (time.perf_counter() - t0) * 1e3

print("shape", tuple(y.shape), "dtype", y.dtype, "ms", round(ms, 3))
print("min/max", float(y.min()), float(y.max()))
print("sample", y[0, 0, :4])
assert y.shape == (1, 4, 8)
assert torch.isfinite(y).all()
assert not torch.allclose(y, x)

TensorFlow example using a tiny CNN

TensorFlow’s tf.keras.utils.set_random_seed() sets Python, NumPy, and backend seeds together, and tf.config.experimental.enable_op_determinism() requests deterministic op behavior on the same hardware. Keras Conv2D expects channels-last by default, so the input here is [1, 8, 8, 1].

import time, tensorflow as tf, numpy as np
seed = 1234
tf.keras.utils.set_random_seed(seed)
tf.config.experimental.enable_op_determinism()

x = np.random.default_rng(seed).standard_normal((1, 8, 8, 1)).astype("float32")

m = tf.keras.Sequential([
    tf.keras.layers.Conv2D(2, 3, activation="relu", input_shape=(8, 8, 1)),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(3),
], name="tiny_cnn")

_ = m(x, training=False)  # build
t0 = time.perf_counter()
y = m(x, training=False)
ms = (time.perf_counter() - t0) * 1e3
y_np = y.numpy()

print("shape", y_np.shape, "dtype", y_np.dtype, "ms", round(ms, 3))
print("min/max", float(y_np.min()), float(y_np.max()))
print("sample", y_np[0, :3])
assert y_np.shape == (1, 3)
assert np.isfinite(y_np).all()
assert np.abs(y_np).max() < 1e6

JAX example using a tiny MLP

JAX random generation is explicitly key-based rather than stateful, which makes the dataflow around randomness unusually clear and deterministic. The example below also uses jax.jit; the first call warms up the compiled path and the timed call measures the steady-state single forward pass.

import time, jax, jax.numpy as jnp
seed = 1234
k0, k1, k2, k3 = jax.random.split(jax.random.key(seed), 4)

W1 = jax.random.normal(k0, (4, 8), dtype=jnp.float32)
b1 = jnp.zeros((8,), dtype=jnp.float32)
W2 = jax.random.normal(k1, (8, 3), dtype=jnp.float32)
b2 = jnp.zeros((3,), dtype=jnp.float32)
x = jax.random.normal(k2, (2, 4), dtype=jnp.float32)

def f(x, p):
    h = jnp.maximum(x @ p[0] + p[1], 0.0)
    return h @ p[2] + p[3]

f_jit = jax.jit(f)
_ = f_jit(x, (W1, b1, W2, b2)).block_until_ready()  # compile
t0 = time.perf_counter()
y = f_jit(x, (W1, b1, W2, b2)).block_until_ready()
ms = (time.perf_counter() - t0) * 1e3

print("shape", y.shape, "dtype", y.dtype, "ms", round(ms, 3))
print("min/max", float(y.min()), float(y.max()))
print("sample", y[0, :3])
assert y.shape == (2, 3)
assert jnp.isfinite(y).all()
assert jnp.abs(y).max() < 1e6

ONNX Runtime example using a tiny MLP defined with ONNX helper APIs

ONNX Runtime’s main execution entry point is InferenceSession, and the ONNX Python helper APIs provide make_node, make_tensor_value_info, and make_model for building a minimal graph. ONNX Runtime’s CPU package is onnxruntime; at runtime you can fix execution to the CPU provider and reduce timing variance with sequential execution and a small thread count.

import time, numpy as np, onnx, onnxruntime as ort
from onnx import helper, TensorProto, numpy_helper

seed = 1234
rng = np.random.default_rng(seed)
x = rng.standard_normal((2, 4), dtype=np.float32)
W = rng.standard_normal((4, 3), dtype=np.float32)
b = rng.standard_normal((3,), dtype=np.float32)

X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 4])
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3])
init = [numpy_helper.from_array(W, "W"), numpy_helper.from_array(b, "b")]
nodes = [
    helper.make_node("MatMul", ["X", "W"], ["Z"]),
    helper.make_node("Add", ["Z", "b"], ["Y"]),
]
graph = helper.make_graph(nodes, "tiny_mlp", [X], [Y], init)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])

so = ort.SessionOptions()
so.intra_op_num_threads = 1
so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
sess = ort.InferenceSession(model.SerializeToString(), so,
                            providers=["CPUExecutionProvider"])

_ = sess.run(["Y"], {"X": x})  # warm-up
t0 = time.perf_counter()
y = sess.run(["Y"], {"X": x})[0]
ms = (time.perf_counter() - t0) * 1e3

print("shape", y.shape, "dtype", y.dtype, "ms", round(ms, 3))
print("min/max", float(y.min()), float(y.max()))
print("sample", y[0, :3])
assert y.shape == (2, 3)
assert np.isfinite(y).all()
assert np.abs(y).max() < 1e6

Validation, reproducibility, and common pitfalls

The most useful validation checks are intentionally simple. First, check the exact output shape against the chosen contract. Second, check the exact dtype you intended to run; float64 can sneak in if inputs come from NumPy defaults rather than explicitly cast arrays. Third, check all values are finite. Fourth, verify the output is not degenerate: for example, not all zeros, not exactly equal to the input in a non-identity model, and not absurdly large in magnitude. Fifth, rerun the same fixture in the same environment and compare either exact equality or a small absolute tolerance depending on the framework and device. These are small checks, but together they catch most wiring errors in early fixture work. PyTorch and TensorFlow both document determinism controls; JAX’s explicit keys make repeated-input/repeated-output reasoning straightforward; ONNX Runtime’s execution path is explicit through session.run().

For reproducibility, the first rule is to pin the environment: same framework version, same Python version, same provider/back-end, same hardware target, and ideally same operating system. PyTorch explicitly warns that complete reproducibility is not guaranteed across releases, commits, platforms, or CPU/GPU executions, even with the same seeds. TensorFlow similarly notes that deterministic execution can be slower, and its Keras seeding helper notes limitations in distributed communication and some non-deterministic cuDNN paths. JAX’s key model makes randomness explicit, but you still need to keep the same code path and device behavior. ONNX Runtime requires you to think clearly about providers and threading, because those are part of the execution contract.

The most common pitfalls are predictable. A channel-order mismatch is the classic CNN error: PyTorch Conv2d is channel-first, TensorFlow Conv2D is channel-last by default. A transformer shape mismatch often comes from forgetting batch_first=True in PyTorch or misunderstanding (B, T, dim) expectations in TensorFlow attention APIs. A timing mistake happens when you benchmark the first JAX jit call or include ONNX session construction in your “one-pass” timing number. A determinism mistake happens when dropout is left enabled, the model remains in training mode, or only one RNG is seeded while NumPy or another library is still producing different inputs. Another subtle issue is assuming CPU and GPU outputs should be bit-identical across frameworks; PyTorch’s reproducibility note is a strong warning against that assumption.

A good operational pattern is to store a tiny artifact alongside each fixture run: the framework name and version, device/provider, seed, model hyperparameters, input tensor checksum, output tensor checksum, and timed milliseconds. That lets you distinguish “math changed” from “environment changed” later. If you are using the fixture as a parity gate for a future larger runtime, this tiny metadata record is often as valuable as the output values themselves.

Framework comparison and troubleshooting

The table below mixes facts from the official docs with a small amount of engineering judgment. Code size is counted directly from the snippets in this report. Runtime overhead is qualitative: it reflects whether the framework is primarily eager, JIT-compiled, or session-based, and whether a one-shot tiny fixture typically pays setup cost before steady-state execution. Device-support cells are anchored to the official install or API docs.

A short troubleshooting checklist:

• Run-to-run outputs differ on the same machine: confirm all relevant RNGs are seeded, dropout is disabled, deterministic mode is enabled where applicable, and the model is in inference/eval mode.
• CNN fixture throws shape errors: check whether the framework expects channel-first or channel-last tensors.
• Transformer fixture throws dimension errors: verify (batch, seq, feature) in PyTorch with batch_first=True, and (B, T, dim) in TensorFlow attention APIs.
• JAX timing looks unexpectedly slow: you probably measured compilation; warm up once before timing the actual pass.
• ONNX Runtime is not using the accelerator you expected: specify the execution provider explicitly instead of relying on defaults for non-CPU targets.
• TensorFlow deterministic mode slows the fixture down: that is expected; the docs explicitly note a performance trade-off.
• CPU and GPU outputs are close but not bit-identical: treat that as normal unless your test contract explicitly requires same-device exact reproducibility. PyTorch warns against assuming cross-platform or CPU/GPU exact reproducibility.

The practical recommendation is simple: make the MLP fixture your universal first gate, the Transformer block your LLM-like second gate, and keep the ONNX Runtime fixture if your long-term path includes deployment validation or a future .NET runtime surface. That combination gives you the best balance of speed, determinism, and diagnostic value.