[Docs] Update website for V2

Author: dbrll

docs/index.md

@@ -10,9 +10,13 @@ hide:
   <h1 style="margin: 0; font-size: 4.5rem; font-weight: 500; color: black;">Sheaf</h1>
 </div>
 
-<h2>A Functional Language for Differentiable Computation</h2>
+<h2 style="max-width: 36rem; margin-left: auto; margin-right: auto; margin-top: 1rem;">A functional language for differentiable computation</h2>
 
-<p>Inspired by Clojure. Designed from the ground up for machine learning. Compiles to native GPU code.</p>
+<p  style="text-align: center; size: 1rem; max-width: 40rem; margin-left: auto; margin-right: auto; margin-top: 1rem;"  markdown="1">
+
+Sheaf brings Clojure’s code-as-data to machine learning, with models as inspectable, composable, and compiled data structures.
+
+</p>
 
 </div>
 
@@ -22,9 +26,9 @@ hide:
 
 ## For ML Researchers
 
-- **Single binary framework** — Models train and run on GPU with no environment to configure
-- **Functional parameters** — Models are data structures that can be inspected, transformed, and composed
-- **Runtime Observability** — Guards and traces expose tensor stability and failure modes
+- **No classes, no boilerplate** — Write math, not plumbing
+- **Runtime Observability** — Catch NaN, trace shapes and profile performance without code changes
+- **Single binary framework** — One executable, no dependencies. Train and run on GPU out of the box
 
 </div>
 
@@ -45,11 +49,9 @@ hide:
 <div class="flex-section" markdown="1">
 <div markdown="1">
 
-Neural networks in Sheaf are written as mathematical transformations.
+In Sheaf, a neural network is a composition of mathematical functions over a parameter tree.
 
-Layers, activations, and parameter bindings form explicit data flow, without imperative state management.
-
-Because the language is functionally pure, compilation and differentiation require no decorators or annotations. `value-and-grad` differentiates any pure function; the JIT resolves shapes and compiles to GPU code automatically.
+Sheaf is purely functional, so differentiation and compilation require no annotations. Any pure function can be differentiated with value-and-grad and is automatically compiled to GPU code.
 
 </div>
 <div markdown="1">
@@ -98,26 +100,23 @@ Because the language is functionally pure, compilation and differentiation requi
 <div class="flex-section" markdown="1">
 <div markdown="1">
 
-In Sheaf, model parameters are nested dictionaries.
-
-There are no module classes, or registration, or parameter groups. Structural operations like pruning, freezing, or weight sharing are expressed as regular data transformations.
+Because models are data, Sheaf requires no module classes, registration, or parameter groups. Even structural operations like pruning, freezing, or weight sharing are expressed as regular data transformations.
 
-Compile-time macros generate architecture variants from a single template. The same primitives extend to neuro-symbolic pipelines where logic and learning are jointly differentiable.
+Sheaf brings compile-time macros to the computation graph itself, generating architecture variants from a single template.
 
 </div>
 
 <div markdown="1">
 
 ```clojure
-;; Freeze a layer: zero out its gradients
-(defn freeze [grad layer-key]
-  (assoc grad layer-key
-    (tree-map (fn [g] (zeros (shape g)))
-              (get grad layer-key))))
+;; Grow a model: add a layer at runtime
+(defn append-layer [params new-layer]
+  (assoc params :layers
+    (append (get params :layers) new-layer)))
 
-;; Apply weight decay to all parameters at once
-(defn weight-decay [params rate]
-  (tree-map (fn [w] (* w (- 1.0 rate))) params))
+;; Swap the output head for a different task
+(defn hot-swap-head [model task-id heads]
+  (assoc model :head (get heads task-id)))
 
 ```
 
@@ -129,9 +128,9 @@ Compile-time macros generate architecture variants from a single template. The s
 <div class="flex-section" markdown="1">
 <div markdown="1">
 
-Every function call, tensor shape, and numerical statistic is observable at runtime.
+In Sheaf, every function call, tensor shape, and numerical statistic is observable at runtime.
 
-Three modes expose different aspects of execution: a tracer reconstructs the full call hierarchy with tensor statistics, guards halt on numerical invariants like NaN or range violations, and a profiler attributes wall time to each function in the call tree.
+A tracer logs the full call hierarchy with tensor statistics. Guards halt execution on numerical invariants like NaN or range violations. A profiler attributes wall time to each function in the call tree.
 
 </div>
 
@@ -147,9 +146,7 @@ Three modes expose different aspects of execution: a tracer reconstructs the ful
     │ │ ├─ [sigmoid] f32[4x1] [min:-5.48e-2 max:1.18e0] (16B)
     │ │ └─ ← f32[4x1] [min:4.86e-1 max:7.66e-1] (16B) (1.8μs)
     │ └─ ← f32[4x1] [min:4.86e-1 max:7.66e-1] (16B) (0.0μs)
-    │ ├─ [sgd-step] dict(keys:["l1", "l2"]), dict(keys:["l1", "l2"]), 0.700000
-    │ └─ ← dict(keys:["l1", "l2"]) (8.9μs)
-    └─ ← dict(keys:["loss", "p"]) (0.0μs)
+    ...
     ```
 
 === "Guards"
@@ -176,11 +173,11 @@ Three modes expose different aspects of execution: a tracer reconstructs the ful
 === "Profiler"
 
     ```css
-    Profiler: 5.63s wall
+    Profiler: 3.63s wall
 
       Function                          Calls      Total       Self   Avg/call
       ------------------------------------------------------------------------
-      gpt-forward                         100      3.72s      3.72s    37.23ms
+      gpt-forward                         100      1.72s      3.72s    37.23ms
       reshape                             301   900.57ms   900.57ms     2.99ms
       choice                              100   622.85ms   622.85ms     6.23ms
       softmax                             100   158.67ms   158.67ms     1.59ms
@@ -191,11 +188,11 @@ Three modes expose different aspects of execution: a tracer reconstructs the ful
 
       Call tree:
 
-      ├── generate (5.58s, 1 call)
-      │   ├── reduce (5.58s, 1 call)
-      │   │   └── <lambda> (5.58s, 101 calls)
-      │   │       ├── generate-token (5.56s, 100 calls)
-      │   │       │   ├── gpt-forward (3.72s, 100 calls)
+      ├── generate (3.58s, 1 call)
+      │   ├── reduce (3.58s, 1 call)
+      │   │   └── <lambda> (3.58s, 101 calls)
+      │   │       ├── generate-token (3.56s, 100 calls)
+      │   │       │   ├── gpt-forward (1.72s, 100 calls)
       │   │       │   ├── reshape (900.16ms, 100 calls)
       │   │       │   ├── choice (622.85ms, 100 calls)
       │   │       │   ├── softmax (158.67ms, 100 calls)
@@ -213,7 +210,7 @@ Three modes expose different aspects of execution: a tracer reconstructs the ful
 <div class="flex-section">
 <div markdown="1">
 
-A complete GPT-2 124M implementation in Sheaf is 1,908 tokens, while the equivalent PyTorch is 7,486. Sheaf's uniform syntax keeps the code concise and close to the math.
+A complete GPT-2 124M implementation in Sheaf is 1,908 tokens, while the equivalent PyTorch is 7,486. Sheaf's uniform syntax keeps the code concise and unambiguous.
 <br/><br/>
 
 Context usage counts GPT-4 tokens (tiktoken) across model, training, and sampling code. Deploy size is the minimal runtime required to train and run a model on a CUDA GPU.
@@ -260,29 +257,21 @@ Context usage counts GPT-4 tokens (tiktoken) across model, training, and samplin
 Sheaf is written in Rust. The complete runtime with GPU backends ships
 as a single 4 MB executable.
 
-Sheaf code is JIT-compiled to optimized GPU kernels for CUDA and Metal through StableHLO and IREE.
-
 The compiler toolchain is downloaded on first use and is not required to run a compiled model.
 
 </div>
 
 <div markdown="1">
 
 ```
-$ du -h * # Standalone, self-contained deployment
+# Standalone, self-contained deployment
+$ du -h *
 128K	__sheaf__                 # compiled model
 3.2M	data
 4.0K	model.shf
-164M	out-shakespeare-char
-3.6M	sheaf                     # runtime
+164M	out-weights
+3.8M	sheaf                     # runtime
 4.0K	train.shf
-
-$ ./sheaf train.shf
-Loading model...
-Loaded: 6 layers, 384 dim, 65 vocab
-Training data: 1115394 tokens
-Training for 100 steps (batch_size=4 block_size=256)...
-Step 100 | Loss: 2.4573
 ```
 
 </div>

docs/key-concepts.md

@@ -77,7 +77,7 @@ The reverse-mode pass then walks these bindings backward, accumulating adjoint c
 | Category    | Operations                                                  |
 | ----------- | ----------------------------------------------------------- |
 | Arithmetic  | `+`, `-`, `*`, `/`, `**`, `neg`                             |
-| Matrix      | `@` (matmul), `transpose`                                   |
+| Matrix      | `@` (matmul), `einsum`, `transpose`                         |
 | Shape       | `reshape`, `swapaxes`, `slice`                              |
 | Activations | `relu`, `gelu`, `sigmoid`, `tanh`, `softmax`, `log-softmax` |
 | Reductions  | `sum`, `mean`, `var`                                        |
@@ -129,13 +129,7 @@ Subsequent calls with the same argument shapes dispatch directly to the cached V
 
 ### Backend Selection
 
-Sheaf selects the compilation backend based on available hardware:
-
-| Platform | Primary Backend | Fallback    |
-| -------- | --------------- | ----------- |
-| Linux    | CUDA            | Vulkan, CPU |
-| macOS    | Metal           | CPU         |
-| Other    | CPU             | -           |
+Sheaf probes available hardware in preference order: CUDA, Metal, Vulkan, CPU. The first available backend is used.
 
 You can override this with the `--device` flag:
 
@@ -246,5 +240,5 @@ Each argument position holds up to 8 cached buffers with most-recently-used orde
 
 - **Shape specialization.** Functions are compiled for the tensor shapes seen at the first call. A different shape (e.g., a different batch size) triggers automatic recompilation and caching. This is the same tracing model as `jax.jit`.
 - **Control flow.** `if` expressions inside compiled functions cause fallback to the interpreter. Conditional logic must live outside the hot path, or use `where`, which compiles to `stablehlo.select`.
-- **Metal backend.** IREE's Metal/SPIR-V backend is less optimized than CUDA. Some operations are slower on Apple GPUs than on equivalent NVIDIA hardware.
+- **Backend maturity.** IREE backend performance varies by hardware. Benchmarks on your target device are recommended.
 - **Memory management.** No automatic model sharding or gradient checkpointing. Large models may exceed device memory.

docs/roadmap.md

@@ -5,26 +5,52 @@ hide:
 
 # Roadmap
 
-## Version 2.1 (planned features)
+These are the planned features for the next releases of Sheaf.
 
-### Performance
+### Language
 
-- **KV cache**: Sheaf 2.0 currently doesn't have a KV cache, which is essential for competitive inference at scale. Without it, each generated token recomputes full O(n^2) attention over the entire context. With it, each step costs O(n) after the first pass.
+- **`def` for global constants**: Immutable top-level bindings, as in Clojure. While not required per-se, it eliminates the need to pass configuration dictionaries everywhere.
 
-- **Batch generation mode**: Currently, autoregressive generation calls the compiled model once per token. Batch mode would compile the full generation loop into a single dispatch, returning all tokens at once.
+- **`loop` / `recur`**: Explicit tail-recursive loops. Sheaf uses `repeat`, but `loop` and `recur` are more natural for someone coming from Clojure.
 
-### Distribution
+- **`reverse` / `flip`**: Reverse a tensor along an axis. Currently requires manual index construction.
 
-- **NCCL all-reduce**: Multi-GPU training via NCCL collective operations, for data-parallel training across multiple devices.
+- **`stack`**: Combine multiple tensors into a new dimension.
 
-### Developer experience
+- **`inc` / `dec`**: Small convenience to increment and decrement, as in Clojure. Currently `(+ var 1)` and `(- var 1)`.
 
-- **Jupyter integration**: A Sheaf kernel for Jupyter, allowing interactive notebook workflows with inline tensor visualization and training loops.
+- **`argmax` returns integers**: `argmax` and `argmin` currently return floats. They will return integer tensors for direct use as indices.
 
-- **`:trace` and `:blame` in the REPL**: These observability modes were tied to V1 semantics and temporarily removed. They will be re-introduced with behavior adapted to the V2 execution model.
+### Macros
+
+- **Enriched `defmacro`**: `range` and `reduce` available at compile time, enabling macros that generate architecture variants from a single template.
 
 ### Operations
 
 - **Convolution primitives**: `conv1d` and `conv2d` via `stablehlo.convolution`, exposed through the standard library.
 
-- **`vmap` on dictionaries**: `vmap` currently only accepts tensor arguments. PyTree support (automatic flattening/unflattening of dicts) will be added to match the behavior of `value-and-grad`.
+- **`vmap` on dictionaries**: `vmap` currently only accepts tensor arguments. PyTree support (automatic flattening/unflattening of dicts) should be added to match the behavior of `value-and-grad`.
+
+### Autodiff
+
+- **Gradient checkpointing**: Recompute forward activations during the backward pass instead of storing them all. This will reduce memory usage for deep models (the GPT-2 124M training currently uses 13 GB).
+
+- **Scalar parameters in `value-and-grad`**: Float scalars in parameter dictionaries (e.g., `{:w 5.0}`) will produce correct gradients. Currently requires wrapping in a 1-element tensor.
+
+### Performance
+
+- **KV cache helpers**: A `cached-attention` stdlib function to simplify implementing KV cache in transformer models.
+
+- **Batch generation mode**: Compile the full autoregressive generation loop into a single dispatch, returning all tokens at once.
+
+### Developer experience
+
+- **Error call stack propagation**: When an error occurs inside a stdlib function, the error message will show the user's call site, not the stdlib internals.
+
+- **Jupyter integration**: A Sheaf kernel for Jupyter, allowing interactive notebook workflows with inline tensor visualization and training loops.
+
+- **`:trace` and `:blame` in the REPL**: These observability modes were tied to V1 semantics and temporarily removed. They will be re-introduced with behavior adapted to the V2 execution model.
+
+### Distribution
+
+- **NCCL all-reduce**: Multi-GPU training via NCCL collective operations, for data-parallel training across multiple devices.