Agent Guidelines

Follow Google Python Style Guide, with JAX patterns in this document taking precedence.

Virtual environment: Use uv for environment management. Run commands with uv run.

Design Decisions

Match the paper. Implementation should follow equations in notes/Whitelam - 2026 - Generative Thermodynamic Computing.pdf.
DRY (Don't Repeat Yourself). Consolidate duplicated implementations into a single source of truth.
Let it crash. Avoid defensive parameter checks; assume correct wiring and let errors surface.
Uncertain correctness. For uncertain behavior, refer to notes/ or ask the user.
Julia-style defaults. Put defaults in function signature: def foo(x, y=10): not def foo(x, y=None): y = y or 10.
No unnecessary intermediate variables. Return directly: return expr not result = expr; return result.
No unused imports/variables. Remove any defined but unreferenced code.

Use jax.lax.scan over Python loops.
Use jax.vmap for batch operations (e.g., multiple denoising trajectories).
No jax.block_until_ready in hot paths. It breaks XLA fusion.
Configurable dtype via config.dtype_name (float32 or float64).