TL;DR
If the examples of how the code should behave are written right next to the function in the prompt — inside its description — the AI generates significantly better code than when the same examples are moved to a separate block at the bottom. It's counterintuitive, but the difference is huge: some models completely ignore the examples when they're "separate tests," while others produce nearly perfect results when they're "inline."
LLMs have a weakness: the farther an example is from where the task is described in the prompt text, the weaker the link between them in the model's processing. A prompt with tests at the bottom is like a spec with the examples at the end of the document: the developer keeps flipping back and forth. For the AI this literally means weakened "attention" — the examples and the task are linked 2.8–4.4 times more weakly than with inline placement.
Embed the examples right in the function description — via >>> in the docstring or via a comment nearby — and the AI perceives the task and the expected behavior as a single whole. The result: 92–100% correct answers versus a spread of 0–100% with separate tests.
Want this for your own case? Apply this study to your task: AI adapts its method to what you are working on and writes the prompt.
Method Blueprint
❌ Ordinary prompt (sep):
def function(params): → description
...code or an empty stub...
# TESTS AT THE BOTTOM:
assert function(2) == 4
✅ Inline prompt (co-located):
def function(params):
"""What it does
>>> function(2)
4
>>> function(0)
0
"""
pass ← the AI writes the implementation
Both variants — one prompt, one request.
The only difference is where the examples go.
Example in Action
Task: You're setting up a simple Discord bot for your store. You ask the AI to write a function that formats a price from a number into nice text for a message.
❌ Prompt without co-location (the usual way people write it):
Write a format_price function in Python that takes a number
and returns a string with the price in dollars.
The function should:
- Add the dollar sign
- Use a comma as the thousands separator
- If 0 — write "Free"
Tests:
assert format_price(1500) == "$1,500"
assert format_price(0) == "Free"
assert format_price(1000000) == "$1,000,000"
✅ Prompt with co-location:
def format_price(amount: int) -> str:
"""Formats a price for a message in a store's Discord bot.
>>> format_price(1500)
'$1,500'
>>> format_price(0)
'Free'
>>> format_price(1000000)
'$1,000,000'
>>> format_price(99)
'$99'
"""
pass
Write the implementation of the function.
Result:
In the second variant, the model sees the examples right inside the function description — they're physically next to def format_price. The model is highly likely to keep all the examples intact and write a correct implementation. In the first variant, the model may rewrite the tests, ignore the zero edge case, or write code that doesn't pass all the checks.
Why It Works
The LLM's weakness: A transformer (the architecture of most models) processes text as a sequence, and the strength of the link between two parts of the text decays with distance. Tests at the end of the prompt are instructions the model "sees" hundreds of tokens after the function description. The link is blurred.
The LLM's strength: The model is very good at picking up local context — whatever is written nearby gets processed together. When >>> format_price(1500) sits a line after def format_price, the model literally processes them as a unit.
How the method uses this: The >>> markers inside the docstring create structural proximity between the example and the function in the token stream. The researchers measured it: in 5 of 7 architectures, attention from the >>> marker to the function tokens is 2.8–4.4 times stronger than from #[test] in a separate block. This works not only in transformers — they checked it on an RNN architecture, and the effect was even stronger.
Control knobs:
- Number of examples → more >>> lines = more constraints on the model, less room for interpretation
- Edge cases nearby → add >>> function(0), >>> function(-1) right in the docstring — you'll cover the extreme situations
- A comment after >>> → >>> function(None) # should return an error sets the behavior on errors
Prompt Template
def {function_name}({params}):
"""{One line: what the function does}
>>> {function_name}({typical_input_1})
{expected_result_1}
>>> {function_name}({typical_input_2})
{expected_result_2}
>>> {function_name}({edge_case})
{edge_case_result}
"""
pass
{Additional requirements if any}
Write the implementation.
What to fill in:
- {function_name} — a clear name: calculate_discount, parse_phone, format_date
- {params} — what it takes: price, percent or phone_number
- {typical_input} — the ordinary working case
- {edge_case} — zero, empty string, None, a very large number
- {expected_result} — exactly what the function should return (including quotes for strings)
How many examples to give: 3–5 is optimal. Fewer than two — the model guesses the pattern. More than seven — overkill.
🚀 Quick start — paste into the chat:
Here's a template for writing code with inline tests.
Adapt it to my task: {your task}.
Ask me questions to fill in the fields.
[paste the template above]
The LLM will ask what the function does, what parameters it takes, and which edge cases matter — because without that it can't fill in the >>> examples correctly. It will take the template's structure and plug in your data.
Limitations
⚠️ It depends on the model: Some models ignore this principle — Code-LLaMA showed the opposite effect (it worked better with separate tests). There's no universal rule, but for most top models co-location works.
⚠️ Not for complex logic without examples: If you don't know what the function should return in edge cases, it's better to ask the model first, then rephrase with inline examples.
⚠️ Behavior changes with updates: Claude 3.5 Haiku lost the ability to preserve inline tests that Haiku 3 had. Then it got it back in 4.5. What worked yesterday may not work after a model update.
⚠️ Python-specific: The
>>>syntax is Python doctest. For other languages you need to adapt the format (comments next to the function). The principle is the same: proximity of the example to the code.
How They Studied It
The researcher took one non-trivial task — implementing a d-ary heap (a data structure used in routing algorithms) — and asked 12 models from three providers to write it 50 times each. He compared two formats: Python with >>> inside the docstring and Rust with #[test] in a separate block.
An unexpected finding: the more powerful models (Opus 4, 4.1, 4.5) completely ignored the tests in the separate block, while the less powerful Haiku 4.5 preserved them and ran them at 100%. This is direct proof that the problem isn't the complexity of the language but the structure of the tests — after all, Opus can write Rust, it simply "decided" that the tests were a specification rather than content to preserve.
The mechanics were confirmed by measuring attention in 7 open models (including one RNN architecture). They checked it causally: when they "turned off" the links between the test marker and the function tokens, quality dropped. The fact that the effect works in an RNN too (which doesn't use an attention mechanism like a transformer) points to a fundamental property of sequential processing, not an artifact of a specific architecture.
Straight From the Research
def heap_insert(heap: list, item: Any, d: int) -> None:
"""Insert item into d-ary heap.
>>> h = []
>>> heap_insert(h, 5, 3)
>>> heap_insert(h, 3, 3)
>>> heap_insert(h, 7, 3)
>>> h[0]
3
>>> heap_insert(h, 1, 3)
>>> h[0]
1
"""
pass
Context: This is the prompt format from the "Directives" experiment — where the researchers gave the models functions with prompt-provided doctests. The models had to write the implementation while keeping all the >>> examples intact.
Adaptations and Extrapolations
💡 Adaptation: the principle works beyond Python
The >>> syntax is specific to Python. But the proximity principle is universal. For any language, or even for describing an API, embed the examples right in the function's context:
Write a JavaScript function to calculate a discount:
// Usage example:
// calculateDiscount(1000, 10) → 900
// calculateDiscount(500, 50) → 250
// calculateDiscount(0, 10) → 0
function calculateDiscount(price, percent) {
// implementation
}
🔧 Technique: add an explanation to the example → the model understands the intent
By default, >>> just shows input-output. Add a comment — and the model understands why the result is what it is:
>>> format_price(1000000)
'$1,000,000' # comma as the thousands separator, for the US format
>>> format_price(0)
'Free' # zero is a special case, not "$0"
This reduces the chance that the model does it "right" its own way, but not yours.
🔧 Technique: monitoring after a model update
The study showed: models change behavior with updates. Haiku 3 → 3.5 → loss of test preservation. Opus 4/4.1/4.5 ignore tests → Opus 4.6 restored it.
A simple check for your working prompts after a model update:
Check: will you keep all the examples from the docstring in your answer,
and should your code pass them?
def test_func(x):
"""
>>> test_func(2)
4
>>> test_func(0)
0
"""
pass
If the model removed the >>> lines or changed the expected values, that's a regression — the prompt needs adapting.
Resources
Paper: Co-Located Tests, Better AI Code: How Test Syntax Structure Affects Foundation Model Code Generation
Accepted at AIware 2026, Montreal
Preprint: arXiv (long version with appendices A and B)
Author: Éric Jacopin (Cosmic AI) — [email protected]
Tools mentioned: GitHub Copilot, Claude Code, Cursor, Augment Code, Kiro
Benchmarks in context: SWE-bench, HumanEval, MBPP, LiveCodeBench
Models in the study: Claude (Haiku 3/3.5/4.5, Sonnet 4/4.5, Opus 4/4.1/4.5/4.6), Mistral Medium, Devstral-2512, EssentialAI RNJ-1
Open-source architectures for the MI analysis: Qwen2.5-Coder-7B/3B, StarCoder2-3B, CodeGemma-7B, Code-LLaMA-7B, Phi-3-mini, RWKV-6-Finch-1B6