hhlo 0.4.0.0 → 0.5.0.0

raw patch · 21 files changed

+2209/−299 lines, 21 filesnew-component:exe:example-autograd-basicnew-component:exe:example-autograd-compositenew-component:exe:example-autograd-linearPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +22/−1
• README.md +270/−294
• examples/34-autograd-basic.hs +54/−0
• examples/35-autograd-linear.hs +62/−0
• examples/36-autograd-composite.hs +67/−0
• hhlo.cabal +47/−1
• src/HHLO/Autograd.hs +29/−0
• src/HHLO/Autograd/Core.hs +294/−0
• src/HHLO/Autograd/Grad.hs +195/−0
• src/HHLO/Autograd/Rules.hs +599/−0
• src/HHLO/EDSL/Ops.hs +206/−0
• src/HHLO/IR/Pretty.hs +2/−2
• test/Main.hs +6/−1
• test/Test/Autograd/Grad.hs +47/−0
• test/Test/Autograd/Rules.hs +49/−0
• test/Test/EDSL/Ops.hs +71/−0
• test/Test/Runtime/EndToEndAutograd.hs +74/−0
• test/Test/Runtime/EndToEndDataMovement.hs +14/−0
• test/Test/Runtime/EndToEndMatmul.hs +41/−0
• test/Test/Runtime/EndToEndReductions.hs +27/−0
• test/Test/Runtime/EndToEndShape.hs +33/−0

@@ -47,7 +47,7 @@ * New comparison wrappers: `equal`, `notEqual`, `greaterThan`, `lessThanOrEqual`, `greaterThanOrEqual`. * Test count: 141 CPU tests + 6 GPU integration tests. -## 0.4.0.0 -- 2026-04-20+## 0.4.0.0 -- 2026-04-26  **BREAKING**: `HostType 'Bool` changed from `Bool` to `Word8` to match PJRT's PRED buffer transfer semantics.@@ -66,3 +66,24 @@ * Boolean logic ops: `logicalAnd`, `logicalOr`, `logicalNot`. * New dependency: `directory` (for plugin-path discovery in `withCPU`/`withGPU`). * Test count: 155 CPU tests + 6 GPU integration tests.+++## 0.5.0.0 -- 2026-04-27++* **Autograd** — reverse-mode automatic differentiation is now part of HHLO.+  New module `HHLO.Autograd` provides `grad` and `vjp` combinators that+  transform HHLO computation graphs into their gradients, producing new+  StableHLO modules that compile via PJRT. VJP rules cover ~25 ops including+  element-wise arithmetic, matmul, transpose, reshape, broadcast, reduce,+  slice, pad, concatenate, select, and more.+* New convenience ops:+  * `einsum` — Einstein summation via subscript strings (e.g. `"ij,jk->ik"`).+    Parses labels, computes batch/contracting dims, and emits the correct+    `stablehlo.dot_general` + optional `stablehlo.transpose`.+  * `split` — split a tensor into N equal parts along a dimension.+  * `stack` — stack N tensors along a new axis.+  * `productAll`, `productDim` — product reductions (mirrors `sumAll`/`reduceSumDim`).+  * `topK` — return top-K values along a dimension via `sort` + `slice`.+* Bug fix: `stablehlo.sort` now wraps its region in parentheses for PJRT+  v1.16.0 parser compatibility.+* Test count: 181 CPU tests + 6 GPU integration tests.

@@ -1,21 +1,86 @@ # HHLO — Haskell Frontend for StableHLO -HHLO is a Haskell library and runtime for building, compiling, and executing machine learning programs targeting [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned intermediate representation of the [OpenXLA](https://openxla.org/) ecosystem.+HHLO is a Haskell library for building, compiling, and executing machine learning programs that target [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned IR of the [OpenXLA](https://openxla.org/) ecosystem. -Instead of replicating JAX's Python-based tracing infrastructure, HHLO generates StableHLO MLIR text directly from Haskell and compiles it to **CPU or GPU** via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) plugin interface.+It lets you write ML models in pure Haskell with compile-time shape checking, compile them to CPU or GPU via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) C API, and even differentiate them automatically — all without leaving the type system. +```haskell+{-# LANGUAGE DataKinds, TypeApplications #-}+import HHLO.Session+import HHLO.EDSL.Ops+import HHLO.Autograd++-- Define a model, differentiate it, and run it on CPU in 6 lines.+main = withCPU $ \sess -> do+    let f x = sumAll =<< multiply x x+        gradMod = gradModule @'[3] @'F32 f+    compiled <- compile sess gradMod+    result   <- run sess compiled (hostFromList @'[3] @'F32 [1, 2, 3])+    print (hostToList result)   -- [2.0, 4.0, 6.0]+```+ --- -## Design+## Table of Contents -HHLO is structured in five layers:+- [Why HHLO?](#why-hhlo)+- [Design Philosophy](#design-philosophy)+- [Features](#features)+  - [Type-Safe EDSL](#type-safe-edsl)+  - [Convenience Ops](#convenience-ops)+  - [Autograd](#autograd)+  - [Runtime & Hardware](#runtime--hardware)+  - [Control Flow & RNG](#control-flow--rng)+- [Quick Start](#quick-start)+- [Examples](#examples)+- [Installation](#installation)+- [Project Structure](#project-structure)+- [License](#license) +---++## Why HHLO?++Most ML frameworks trace Python code to build computation graphs. HHLO takes a different path: you write StableHLO directly in Haskell.++This means:++- **No Python runtime** — Your model is ordinary Haskell code.+- **Compile-time shape safety** — Matmul mismatches are type errors, not runtime failures.+- **Native autograd** — Reverse-mode differentiation is implemented as a Haskell library, not a C++ backend.+- **True portability** — StableHLO is a standardized, versioned IR; the same Haskell code runs on CPU, NVIDIA GPU, or any future PJRT backend.++---++## Design Philosophy++### Text Emission + PJRT++HHLO emits StableHLO MLIR text and hands it straight to `PJRT_Client_Compile`. This is the same compilation path used by JAX's C++ backend, but without the heavy dependency of building LLVM/MLIR from source.++### Phantom Types++Every tensor carries its shape and dtype as phantom type parameters:++```haskell+Tensor '[2, 3] 'F32   -- 2×3 matrix of Float32+Tensor '[4]    'F64   -- 4-element vector of Float64 ```++Matmul, broadcast, and conv shapes are checked at compile time via type families. If the shapes don't match, GHC tells you before you ever load a PJRT plugin.++### Layered Architecture++HHLO is structured so you can use as much or as little abstraction as you need:++``` ┌─────────────────────────────────────┐-│  Convenience (HHLO.Session)         │  One-liners: withCPU, compile, run+│  Session (HHLO.Session)             │  One-liners: withCPU, compile, run ├─────────────────────────────────────┤-│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, relu, etc.+│  Autograd (HHLO.Autograd)           │  grad, vjp, gradModule — reverse-mode AD ├─────────────────────────────────────┤+│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, einsum, etc.+├─────────────────────────────────────┤ │  IR Builder (HHLO.IR.Builder)       │  Stateful monad for constructing MLIR ├─────────────────────────────────────┤ │  Pretty Printer (HHLO.IR.Pretty)    │  Emits StableHLO MLIR text@@ -24,86 +89,136 @@ └─────────────────────────────────────┘ ``` -**Text Emission + PJRT**+The high-level layers (`Session`, `Autograd`) eliminate PJRT boilerplate for the common case. The low-level layers (`IR.Builder`, `Pretty`, `Runtime`) remain available when you need full control. -The library emits StableHLO MLIR text directly and hands it to `PJRT_Client_Compile`. This is the same path used by JAX's C++ backend and avoids the heavy dependency of building LLVM/MLIR from source.+--- -**Phantom Types**+## Features -Every tensor carries its shape and dtype as phantom type parameters:+### Type-Safe EDSL++The frontend provides 50+ typed ops covering arithmetic, linear algebra, reductions, data movement, and neural network primitives:+ ```haskell-Tensor '[2, 3] 'F32   -- 2×3 matrix of Float32+-- Arithmetic+c <- add a b+d <- multiply a b+e <- matmul a b++-- Non-linear+y <- relu x+y <- sigmoid x+y <- softmax x++-- Reductions+s <- sumAll x                    -- reduce all dims → scalar+v <- reduceSumDim @0 x           -- reduce dim 0++-- Data movement+sliced <- slice x [(0, 2), (1, 3)]   -- extract sub-array+padded <- pad x 0 [(1, 1), (0, 0)]   -- pad with zeros+trans  <- transpose x [1, 0]         -- permute dimensions ```-Matmul, broadcast, and conv shapes are checked at compile time via type families. -**ForeignPtr Finalizers**+Shape mismatches are caught at compile time. A matmul between a `[2,3]` and a `[4,5]` tensor is a type error, not a segfault. -PJRT buffers and executables are managed by `ForeignPtr` finalizers that automatically call `PJRT_Buffer_Destroy` and `PJRT_LoadedExecutable_Destroy` when values are garbage-collected. You can still let references drop out of scope without explicit cleanup.+### Convenience Ops -**Dynamic Output Counts**+Beyond the raw StableHLO surface, HHLO provides higher-level combinators that compose primitive ops into familiar patterns: -The runtime queries the compiled executable for its actual number of outputs via `PJRT_Executable_NumOutputs` instead of guessing or hardcoding a maximum.+| Op | What it does |+|---|---|+| `einsum "ij,jk->ik" a b` | Einstein summation (dispatches to `dotGeneral` + `transpose`) |+| `split dim n t` | Decompose a tensor into `n` equal slices along `dim` |+| `stack dim [t1, t2, ...]` | Concatenate tensors along a new axis `dim` |+| `productAll t` | Reduce all dimensions with `multiply` (like `sumAll` but product) |+| `productDim dims t` | Reduce specific dimensions with `multiply` |+| `topK k t` | Return the top-K elements (sort descending + slice) | -**Async Execution**+These are implemented as pure compositions of existing EDSL ops, so they inherit full autograd support automatically. -`HHLO.Runtime.Async` provides true non-blocking execution: `executeAsync` returns buffer handles immediately, `bufferReady` polls for completion, and `awaitBuffers` blocks until device-side computation finishes.+### Autograd -**Device Enumeration & Selection**+HHLO includes a native reverse-mode automatic differentiation engine that transforms StableHLO computation graphs into their gradients. -`HHLO.Runtime.Device` lets you discover and select specific GPUs at runtime:+**Standalone modules** — produce a reusable `Module`:+ ```haskell-addressableDevices api client        -- list all devices-deviceKind api dev                   -- "cpu" or "NVIDIA GeForce RTX 5090"-defaultGPUDevice api client          -- first non-CPU device+-- f(x) = sum(x²)   =>   grad f(x) = 2x+gradMod :: Module+gradMod = gradModule @'[3] @'F32 $ \x -> do+    sq <- multiply x x+    sumAll sq ``` -**Multi-GPU Inference Scaling**+**In-place combinators** — use inside `buildModule` for composability: -`HHLO.Runtime.Execute` provides `executeReplicas` for running the same compiled model concurrently across multiple GPUs: ```haskell-compileWithOptions api client mlirText-    (defaultCompileOptions { optNumReplicas = numDevs })+buildModule @1 @1 "loss_and_grad" $ \x -> do+    loss <- sumAll =<< multiply x x+    g    <- grad (\y -> sumAll (multiply y y)) x+    returnTuple2 loss g+``` --- Launch independent forward passes on all GPUs-executeReplicas api exec-    [ (gpu0, [bufA0, bufB0])-    , (gpu1, [bufA1, bufB1])-    , ...-    ]+**Vector-Jacobian products** — for non-scalar outputs:++```haskell+-- vjp f x seed = (Df(x))ᵀ · seed+vjpModule @'[3] @'[2] @'F32+    (\x -> do w <- constant @'[2,3] @'F32 1.0; matmul w x) ``` -**Multi-Result Operations**+Supported ops: `add`, `subtract`, `multiply`, `divide`, `negate`, `exponential`, `log`, `sqrt`, `power`, `sine`, `cosine`, `tanh`, `abs`, `maximum`, `minimum`, `reshape`, `transpose`, `broadcast_in_dim`, `reduce` (sum), `dot`, `select`, `slice`, `pad`, `concatenate`, `convert`, and more. -The AST `Operation` type supports multiple results, enabling ops like `stablehlo.rng_bit_generator` and multi-value control flow:+Ops without gradient rules (e.g. `compare`, `floor`, `ceil`, `sort`) safely return zero gradients. Stubs (e.g. `gather`, `scatter`) error explicitly.++### Runtime & Hardware++**CPU & GPU**++The same Haskell code compiles to CPU via `withCPU` or to GPU via `withGPU`:+ ```haskell--- Two-result operation-(newState, output) <- rngBitGenerator state+withCPU $ \sess -> do ...   -- CPU plugin, works out of the box+withGPU $ \sess -> do ...   -- CUDA plugin, requires NVIDIA runtime libs ``` -**Convenience Layer**+**Async Execution** -`HHLO.ModuleBuilder` and `HHLO.Session` provide a high-level API that eliminates PJRT boilerplate for the common case:+`HHLO.Runtime.Async` provides true non-blocking execution:  ```haskell-import HHLO.ModuleBuilder-import HHLO.Session+bufs <- executeAsync api exec inputs+ready <- bufferReady api (head bufs)   -- poll+awaitBuffers api bufs                   -- block until done+``` --- Build + compile + run in four lines-main = withCPU $ \sess -> do-    let modu = buildModule @2 @1 "mul" $ \x y -> multiply x y-    compiled <- compile sess modu-    result <- run sess compiled (hostFromList @'[2] [2.0, 3.0],-                                  hostFromList @'[2] [4.0, 5.0])-    print (hostToList result)   -- [8.0, 15.0]+**Multi-GPU Inference**++Run the same compiled model concurrently across multiple GPUs:++```haskell+compileWithOptions api client mlirText+    (defaultCompileOptions { optNumReplicas = numDevs })++executeReplicas api exec+    [ (gpu0, [bufA0, bufB0])+    , (gpu1, [bufA1, bufB1])+    , ...+    ] ``` -No `FuncArg`, no `natVal`, no `render`, no `toDeviceF32`, no explicit shape lists. The low-level API remains available for expert users who need full control.+**ForeignPtr Finalizers** +PJRT buffers and executables are managed by `ForeignPtr` finalizers. They are automatically destroyed when garbage-collected — no explicit cleanup required.++### Control Flow & RNG+ **Multi-Value Control Flow**  `whileLoop2` / `conditional2` carry multiple typed tensors through loops and conditionals without manual packing:+ ```haskell--- Loop with two accumulators: counter and running sum (resultCounter, resultSum) <- whileLoop2 counter0 sum0     (\c s -> compare c limit "LT")     (\c s -> do@@ -114,294 +229,157 @@  **Random Number Generation** -Three RNG primitives are exposed in the EDSL: ```haskell uniform  <- rngUniform a b      -- uniform in [a, b) normal   <- rngNormal            -- standard normal (mean 0, std 1) (newSt, bits) <- rngBitGenerator state   -- Threefry bit generator ``` -**Extended Math Primitives**--Element-wise ops covering the full HBayesian requirements:-```haskell-y <- sqrt x          -- square root-y <- rsqrt x         -- reciprocal sqrt-y <- sin x           -- sine-y <- cos x           -- cosine-y <- tan x           -- tangent-y <- pow x e         -- element-wise power-y <- log1p x         -- log(1+x)-y <- floor x         -- floor-y <- ceil x          -- ceiling-y <- sigmoid x       -- 1 / (1 + exp(-x))-```--**Shape-Preserving Comparisons**--`compare` and its wrappers return `Tensor s 'Bool` (same shape as inputs), matching StableHLO semantics:-```haskell-mask <- equal x y                -- element-wise equality-mask <- greaterThan x y          -- element-wise >-mask <- lessThanOrEqual x y      -- element-wise <=-```--Convenience ops for scalar manipulation:-```haskell-s <- sumAll x          -- reduce all dimensions to scalar-v <- slice1 vec i      -- extract scalar from 1-D tensor-packed <- pack2 a b    -- pack two scalars into [2]-```- --- -## Installation--### System Requirements--- GHC 9.6+ and Cabal 3.10+-- Linux x86_64 (other platforms supported by PJRT artifacts may work)-- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`)-- `libstdc++` and `libdl` (usually present on Linux)--### From Hackage--HHLO is published on [Hackage](https://hackage.haskell.org/package/hhlo). You can add it directly to your `.cabal` file:--```cabal-build-depends: hhlo >= 0.4-```--Or with `cabal`:--```bash-cabal install hhlo-```--### Download PJRT Plugins+## Quick Start -Run the provided script to download prebuilt PJRT plugins:+### 1. Download PJRT plugins  ```bash ./pjrt_script.sh ``` -This downloads `libpjrt_cpu.so` from the [zml/pjrt-artifacts](https://github.com/zml/pjrt-artifacts) nightly builds into `deps/pjrt/`. If you have an NVIDIA GPU with `nvidia-smi` available, the CUDA plugin is also fetched automatically.+This fetches `libpjrt_cpu.so` into `deps/pjrt/`. If you have an NVIDIA GPU, the CUDA plugin is also downloaded automatically. -### Build the Project+### 2. Build  ```bash cabal build all ``` -This compiles the library, the demo, the examples, and the test suite.-------## Usage--### CPU (works out of the box)+### 3. Run an example  ```bash+# CPU — works out of the box cabal run example-add --flag=examples-cabal test-``` -> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command.--### GPU (requires runtime libraries)--The PJRT CUDA plugin depends on NVIDIA runtime libraries: **cuDNN**, **NCCL**, and **NVSHMEM**. These are commonly available via conda, pip, or system packages.--If you already have them (e.g. via PyTorch or JAX installations), simply run:--```bash-./setup_gpu_env.sh-source ~/.bashrc+# Autograd+cabal run example-autograd-basic --flag=examples ``` -This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly:+### 4. Run tests  ```bash-cabal run example-gpu-add --flag=examples-cabal run example-gpu-matmul-bench --flag=examples-cabal run example-multi-gpu-inference --flag=examples+cabal test                    # 181 CPU tests+cabal test --test-options="-t HHLO+GPU"   # + 6 GPU integration tests ```  --- -## EDSL Quick Start--### Convenience layer (recommended)--```haskell-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE TypeApplications #-}+## Examples -import HHLO.ModuleBuilder-import HHLO.Session+Standalone examples live in `examples/` and cover arithmetic, neural networks, control flow, RNG, and autograd: --- Build a program: c = a + b-program = buildModule @2 @1 "add" $ \a b -> add a b+| # | Command | Description |+|---|---------|-------------|+| 1 | `example-add` | Element-wise `c = a + b` |+| 2 | `example-matmul` | 2×3 @ 3×2 matrix multiply |+| 3 | `example-chain-ops` | `(a + b) * (a - b)` |+| 4 | `example-async` | Async `executeAsync` + `relu` |+| 5 | `example-mlp` | 2-layer MLP |+| 6 | `example-mlp-batched` | Batched MLP |+| 7 | `example-tuple` | Multi-result `func.func` |+| 8 | `example-reduce` | `reduceSum` over all dimensions |+| 9 | `example-softmax` | 1-D and batched 2-D softmax |+| 10 | `example-conv2d` | NHWC conv2d |+| 11 | `example-batch-norm` | Batch norm inference |+| 12 | `example-while` | `whileLoop` count-up |+| 13 | `example-conditional` | `conditional` if-then-else |+| 14 | `example-gather` | `gather` rows from matrix |+| 15 | `example-scatter` | `scatter` replace into vector |+| 16 | `example-slice` | `slice` sub-array extraction |+| 17 | `example-pad` | `pad` with edge/interior padding |+| 18 | `example-dynamic-slice` | `dynamicSlice` runtime indices |+| 19 | `example-sort` | `sort` 1-D ascending |+| 20 | `example-select` | Element-wise ternary `select` |+| 21 | `example-map` | `map` with custom computation |+| 22 | `example-new-ops-smoke-test` | Smoke test for newer ops |+| 23 | `example-resnet` | ResNet-18 toy (8×8 input) |+| 24 | `example-alexnet` | AlexNet toy (16×16 input) |+| 25 | `example-transformer` | Transformer encoder (1×4×16) |+| 26 | `example-unet` | UNet segmentation toy (16×16) |+| 30 | `example-rng-uniform` | `rngUniform` random floats [0,1) |+| 31 | `example-rng-normal` | `rngNormal` standard normal distribution |+| 32 | `example-rng-bit-generator` | `rngBitGenerator` Threefry PRNG |+| 33 | `example-multi-value-loop` | `whileLoop2` with two loop-carried values |+| **34** | **`example-autograd-basic`** | **Gradient of `sum(x²)`** |+| **35** | **`example-autograd-linear`** | **Gradient of linear + MSE loss** |+| **36** | **`example-autograd-composite`** | **Gradient through ReLU + linear + sum** |+| 27 | `example-gpu-add` | GPU smoke test |+| 28 | `example-gpu-matmul-bench` | GPU 4096×4096 benchmark |+| 29 | `example-multi-gpu-inference` | Multi-GPU concurrent matmul | -main = withCPU $ \sess -> do-    compiled <- compile sess program-    result <- run sess compiled-        ( hostFromList @'[2,2] @'F32 [1, 2, 3, 4]-        , hostFromList @'[2,2] @'F32 [5, 6, 7, 8]-        )-    print (hostToList result)   -- [6.0, 8.0, 10.0, 12.0]-```+> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command. -### Low-level API (full control)+### Writing your own model  ```haskell-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE DataKinds, TypeApplications #-} -import HHLO.Core.Types+import HHLO.Session import HHLO.EDSL.Ops-import HHLO.IR.AST (FuncArg(..), TensorType(..))-import HHLO.IR.Builder-import HHLO.IR.Pretty-import qualified Data.Text as T---- Build a program: c = a + b-program :: Module-program = moduleFromBuilder @'[2,2] @'F32 "main"-    [ FuncArg "a" (TensorType [2, 2] F32)-    , FuncArg "b" (TensorType [2, 2] F32)-    ]-    $ do-        a <- arg-        b <- arg-        c <- add a b-        return c--main :: IO ()-main = T.putStrLn (render program)-```--Output:-```mlir-module {-  func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {-      %0 = stablehlo.add %arg0, %arg1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>-      return %0 : tensor<2x2xf32>-  }-}-```+import HHLO.Autograd -### Running the Demo+-- A tiny model: predict y from x via a learned weight.+-- We want the gradient of the squared error.+main = withCPU $ \sess -> do+    let model x = do+            w <- constant @'[1] @'F32 2.0   -- fixed weight for demo+            b <- constant @'[1] @'F32 1.0+            y <- add =<< multiply w x =<< pure b+            tgt <- constant @'[1] @'F32 5.0+            diff <- sub y tgt+            sumAll =<< multiply diff diff -```bash-cabal run hhlo-demo+    let gradMod = gradModule @'[1] @'F32 model+    compiled <- compile sess gradMod+    result   <- run sess compiled (hostFromList @'[1] @'F32 [3.0])+    print (hostToList result)   -- [8.0] ``` -The demo builds a `stablehlo.add` program via the EDSL, compiles it with PJRT CPU, creates F32 input buffers, executes, and reads back the result:--```-=== HHLO End-to-End Demo ===-Loading PJRT CPU plugin...-Plugin loaded.-...-Result: [6.0,8.0,10.0,12.0]-SUCCESS: Results match expected values!-```+--- -### Running Examples+## Installation -Standalone examples are provided in `examples/`:+### System Requirements -| # | Command | Description |-|---|---------|-------------|-| 1 | `cabal run example-add --flag=examples` | Element-wise `c = a + b` |-| 2 | `cabal run example-matmul --flag=examples` | 2×3 @ 3×2 matrix multiply |-| 3 | `cabal run example-chain-ops --flag=examples` | `(a + b) * (a - b)` |-| 4 | `cabal run example-async --flag=examples` | Async `executeAsync` + `relu` |-| 5 | `cabal run example-mlp --flag=examples` | 2-layer MLP |-| 6 | `cabal run example-mlp-batched --flag=examples` | Batched MLP |-| 7 | `cabal run example-tuple --flag=examples` | Multi-result `func.func` |-| 8 | `cabal run example-reduce --flag=examples` | `reduceSum` over all dimensions |-| 9 | `cabal run example-softmax --flag=examples` | 1-D and batched 2-D softmax |-| 10 | `cabal run example-conv2d --flag=examples` | NHWC conv2d |-| 11 | `cabal run example-batch-norm --flag=examples` | Batch norm inference |-| 12 | `cabal run example-while --flag=examples` | `whileLoop` count-up |-| 13 | `cabal run example-conditional --flag=examples` | `conditional` if-then-else |-| 14 | `cabal run example-gather --flag=examples` | `gather` rows from matrix |-| 15 | `cabal run example-scatter --flag=examples` | `scatter` replace into vector |-| 16 | `cabal run example-slice --flag=examples` | `slice` sub-array extraction |-| 17 | `cabal run example-pad --flag=examples` | `pad` with edge/interior padding |-| 18 | `cabal run example-dynamic-slice --flag=examples` | `dynamicSlice` runtime indices |-| 19 | `cabal run example-sort --flag=examples` | `sort` 1-D ascending |-| 20 | `cabal run example-select --flag=examples` | Element-wise ternary `select` |-| 21 | `cabal run example-map --flag=examples` | `map` with custom computation |-| 22 | `cabal run example-new-ops-smoke-test --flag=examples` | Smoke test for newer ops |-| 23 | `cabal run example-resnet --flag=examples` | ResNet-18 toy (8×8 input) |-| 24 | `cabal run example-alexnet --flag=examples` | AlexNet toy (16×16 input) |-| 25 | `cabal run example-transformer --flag=examples` | Transformer encoder (1×4×16) |-| 26 | `cabal run example-unet --flag=examples` | UNet segmentation toy (16×16) |-| 30 | `cabal run example-rng-uniform --flag=examples` | `rngUniform` random floats [0,1) |-| 31 | `cabal run example-rng-normal --flag=examples` | `rngNormal` standard normal distribution |-| 32 | `cabal run example-rng-bit-generator --flag=examples` | `rngBitGenerator` Threefry PRNG |-| 33 | `cabal run example-multi-value-loop --flag=examples` | `whileLoop2` with two loop-carried values |-| **27** | `cabal run example-gpu-add --flag=examples` | **GPU smoke test** |-| **28** | `cabal run example-gpu-matmul-bench --flag=examples` | **GPU 4096×4096 benchmark** |-| **29** | `cabal run example-multi-gpu-inference --flag=examples` | **Multi-GPU concurrent matmul** |+- GHC 9.6+ and Cabal 3.10++- Linux x86_64 (other platforms supported by PJRT artifacts may work)+- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`)+- `libstdc++` and `libdl` (usually present on Linux) ----+### From Hackage -## Tests+```cabal+build-depends: hhlo >= 0.5+``` -### CPU Tests (default)+Or:  ```bash-cabal test+cabal install hhlo ``` -Runs **155 tests** across three tiers:--- **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow.-- **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops.-- **Tier 3 — Runtime integration tests** — Buffer metadata queries, async execution, and error handling.+### GPU Setup -### GPU Tests+The PJRT CUDA plugin depends on **cuDNN**, **NCCL**, and **NVSHMEM**. If you already have them (e.g. via PyTorch or JAX):  ```bash-HHLO_TEST_GPU=1 cabal test+./setup_gpu_env.sh+source ~/.bashrc ``` -Runs the full 155 CPU tests **plus** 6 additional GPU integration tests:--- `EndToEnd.GPU` — GPU availability and device enumeration-- `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU-- `Runtime.AsyncGPU` — Async execution and `bufferReady` polling on GPU-- `Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs--Sample output:-```-HHLO Tests-  EDSL.Ops-    Binary element-wise-      add:                            OK-      ...-  EndToEnd.Arithmetic-    relu:                             OK (0.02s)-    ...-  Runtime.Buffer-    buffer round-trip f32:            OK-  Runtime.Async-    buffer ready after sync execute:  OK (0.02s)-  EndToEnd.GPU-    gpu available:                    OK-  Runtime.BufferGPU-    gpu buffer round-trip f32:        OK-  Runtime.AsyncGPU-    gpu executeAsync + await:         OK-  Runtime.MultiGPU-    execute replicas on all GPUs:     OK+This auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly: -All 147 tests passed (16.27s)+```bash+cabal run example-gpu-add --flag=examples+cabal run example-gpu-matmul-bench --flag=examples ```  ---@@ -419,14 +397,21 @@ │   └── pjrt/               # Downloaded PJRT plugins (.so files) │       └── lib_symlinks/   # Compatibility symlinks for missing library versions ├── doc/                    # Architecture and design documents-├── examples/               # Standalone example programs (01–33)+├── examples/               # Standalone example programs (01–36) ├── src/HHLO/+│   ├── Autograd/           # Reverse-mode automatic differentiation+│   │   ├── Autograd.hs     # Public re-export module+│   │   ├── Core.hs         # BTensor (runtime-typed backward handles)+│   │   ├── Grad.hs         # grad, vjp, gradModule, vjpModule+│   │   └── Rules.hs        # Per-op VJP rules (~25 ops) │   ├── Core/Types.hs       # DType, Shape, HostType type families │   ├── IR/ │   │   ├── AST.hs          # MLIR AST (Operation, Function, Module) │   │   ├── Builder.hs      # Stateful Builder monad + Tensor/Tuple GADTs │   │   └── Pretty.hs       # MLIR text pretty-printer-│   ├── EDSL/Ops.hs         # Type-safe frontend ops (50+ ops)+│   ├── EDSL/Ops.hs         # Type-safe frontend ops (50+ ops + convenience wrappers)+│   ├── ModuleBuilder.hs    # Typeclass-dispatched buildModuleN @M @K+│   ├── Session.hs          # High-level withCPU / withGPU / compile / run API │   └── Runtime/ │       ├── PJRT/ │       │   ├── FFI.hs      # C FFI declarations@@ -434,13 +419,15 @@ │       │   ├── Error.hs    # PJRT error handling │       │   └── Plugin.hs   # Backend-agnostic plugin loading (withPJRT) │       ├── Device.hs       # Device enumeration & selection-│       ├── Compile.hs      # MLIR → PJRT executable-│       ├── Compile.hs      # MLIR → PJRT executable (with `CompileOptions`)+│       ├── Compile.hs      # MLIR → PJRT executable (with CompileOptions) │       ├── Execute.hs      # Synchronous + device-targeted + multi-GPU replica execution │       ├── Async.hs        # Non-blocking execution with PJRT_Event │       └── Buffer.hs       # Host↔device buffer transfers + metadata queries ├── test/ │   ├── Test/+│   │   ├── Autograd/       # Autograd golden & unit tests+│   │   │   ├── Grad.hs+│   │   │   └── Rules.hs │   │   ├── EDSL/Ops.hs │   │   ├── IR/ │   │   │   ├── Builder.hs@@ -450,6 +437,7 @@ │   │   │   └── PrettyControlFlow.hs │   │   ├── Runtime/ │   │   │   ├── EndToEnd*.hs       # CPU E2E test modules+│   │   │   ├── EndToEndAutograd.hs # Numerical autograd verification │   │   │   ├── EndToEndGPU.hs     # GPU availability test │   │   │   ├── Buffer.hs │   │   │   ├── BufferGPU.hs       # GPU buffer integration tests@@ -464,18 +452,6 @@ ├── setup_gpu_env.sh        # Auto-configures LD_LIBRARY_PATH for GPU └── README.md ```-------<!-- ## Architecture Docs -->--<!-- The `doc/` directory contains detailed design documents: -->--<!-- | Document | Contents | -->-<!-- |----------|----------| -->-<!-- | `implementation-design.md` | Four-layer architecture and design decisions | -->-<!-- | `progress-and-remaining-work.md` | Current status, completed features, and backlog | -->-<!-- | `test-suite-documentation.md` | Test catalog and tier descriptions | -->  --- 

@@ -0,0 +1,54 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 34: Basic Autograd — Gradient of sum(x²).+--+-- Demonstrates the simplest use of HHLO.Autograd:+--   f(x) = sum(x * x)   =>   grad f(x) = 2 * x+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-basic++module Main where++import Prelude hiding (maximum)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 34: Basic Autograd ==="+    putStrLn "f(x) = sum(x * x)  =>  grad f(x) = 2 * x"+    putStrLn ""++    -- Define a scalar-valued function.+    let f :: Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            sq <- multiply x x+            sumAll sq++    -- Compute its gradient as a standalone Module.+    let gradMod = gradModule @'[3] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    -- Compile and run.+    compiled <- compile sess gradMod+    let input = hostFromList @'[3] @'F32 [1.0, 2.0, 3.0]+    result <- run sess compiled input :: IO (HostTensor '[3] 'F32)++    putStrLn $ "Input:    " ++ show (hostToList result)+    putStrLn $ "Expected: [2.0, 4.0, 6.0]"++    let resultList :: [Float]+        resultList = hostToList result+    if resultList == [2.0, 4.0, 6.0]+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"

@@ -0,0 +1,62 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 35: Autograd — Gradient of a linear model + MSE loss.+--+-- We define:  y = W * x + b   (with fixed W, b)+-- and loss:   L = (y - target)²+--+-- grad_L w.r.t. x = 2 * W * (y - target)+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-linear++module Main where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 35: Autograd Linear + MSE ==="+    putStrLn "y = W*x + b,  L = (y - target)^2"+    putStrLn "grad_x L = 2 * W * (y - target)"+    putStrLn ""++    -- Fixed parameters: W = 2.0, b = 1.0, target = 5.0+    -- f(x) = (2*x + 1 - 5)^2 = (2*x - 4)^2+    -- grad f(x) = 2 * 2 * (2*x - 4) = 8*x - 16+    -- At x = 3.0: grad = 8*3 - 16 = 8+    let f :: Tensor '[1] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            w    <- constant @'[1] @'F32 2.0+            b    <- constant @'[1] @'F32 1.0+            tgt  <- constant @'[1] @'F32 5.0+            wx   <- multiply w x+            y    <- add wx b+            diff <- sub y tgt+            sq   <- multiply diff diff+            sumAll sq++    let gradMod = gradModule @'[1] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    compiled <- compile sess gradMod+    let input = hostFromList @'[1] @'F32 [3.0]+    result <- run sess compiled input :: IO (HostTensor '[1] 'F32)++    putStrLn $ "Input x:      " ++ show (hostToList result)+    putStrLn $ "Expected grad: [8.0]  (analytical: 8*3 - 16)"++    let resultList :: [Float]+        resultList = hostToList result+    if abs (head resultList - 8.0) < 0.001+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"

@@ -0,0 +1,67 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 36: Autograd through a composite function (ReLU + linear + sum).+--+-- f(x) = sum( relu( w ⊙ x + b ) )+--+-- Demonstrates that autograd works through non-linear ops+-- (max/relu, multiply, add) and reduce (sumAll).+--+-- With w = [2, 3], b = [1, 1], x = [1, 1]:+--   z = w⊙x + b = [3, 4]+--   relu(z) = [3, 4]+--   f(x) = 7+--+-- grad w.r.t. x:+--   d(relu(z))/dz = [1, 1]  (both positive)+--   dx = w ⊙ [1, 1] = [2, 3]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-composite++module Main where++import Prelude hiding (maximum)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 36: Autograd Composite (ReLU + Linear + Sum) ==="+    putStrLn "f(x) = sum( relu( w * x + b ) )"+    putStrLn ""++    let f :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            w    <- constant @'[2] @'F32 1.0+            b    <- constant @'[2] @'F32 1.0+            wx   <- multiply w x+            z    <- add wx b+            zero <- constant @'[2] @'F32 0.0+            reluZ <- maximum z zero+            sumAll reluZ++    let gradMod = gradModule @'[2] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    compiled <- compile sess gradMod+    let input = hostFromList @'[2] @'F32 [1.0, 1.0]+    result <- run sess compiled input :: IO (HostTensor '[2] 'F32)++    putStrLn $ "Input x:       " ++ show (hostToList result)+    putStrLn $ "Expected grad: [1.0, 1.0]  (w=[1,1], b=[1,1])"++    let res :: [Float]+        res = hostToList result+    if all (\v -> abs (v - 1.0) < 0.001) res+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.4.0.0+version:            0.5.0.0 synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU description:     HHLO is a Haskell library and runtime for building, compiling, and executing@@ -63,6 +63,10 @@         HHLO.EDSL.Ops         HHLO.ModuleBuilder         HHLO.Session+        HHLO.Autograd+        HHLO.Autograd.Core+        HHLO.Autograd.Grad+        HHLO.Autograd.Rules         HHLO.Runtime.PJRT.FFI         HHLO.Runtime.PJRT.Types         HHLO.Runtime.PJRT.Error@@ -487,6 +491,8 @@         Test.IR.Builder         Test.EDSL.Ops         Test.ModuleBuilder+        Test.Autograd.Grad+        Test.Autograd.Rules         Test.Runtime.EndToEnd         Test.Runtime.EndToEndArithmetic         Test.Runtime.EndToEndShape@@ -496,6 +502,7 @@         Test.Runtime.EndToEndDataMovement         Test.Runtime.EndToEndMultiValue         Test.Runtime.EndToEndSession+        Test.Runtime.EndToEndAutograd         Test.Runtime.Buffer         Test.Runtime.Async         Test.Runtime.Errors@@ -558,6 +565,45 @@ executable example-multi-value-loop     import:           warnings     main-is:          33-multi-value-loop.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-autograd-basic+    import:           warnings+    main-is:          34-autograd-basic.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-autograd-linear+    import:           warnings+    main-is:          35-autograd-linear.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-autograd-composite+    import:           warnings+    main-is:          36-autograd-composite.hs     build-depends:         base          >= 4.18.2 && < 5,         hhlo,

@@ -0,0 +1,29 @@+-- | Autograd-HHLO: Reverse-mode automatic differentiation for StableHLO.+--+-- This library provides 'grad' and 'vjp' combinators that transform HHLO+-- computation graphs into their gradients, producing new StableHLO modules+-- that compile via PJRT to CPU or GPU.+--+-- The primary safe entry points are 'gradModule' and 'vjpModule', which+-- produce standalone 'Module' values. The in-place 'grad' and 'vjp'+-- combinators can be used inside 'buildModule' for composability.+--+-- Example:+--+-- > import HHLO.ModuleBuilder+-- > import HHLO.Autograd+-- >+-- > -- f(x) = sum(x^2)+-- > -- grad f(x) = 2 * x+-- > gradMod = gradModule @'[3] @'F32 $ \x -> do+-- >     sq <- multiply x x+-- >     sumAll sq+module HHLO.Autograd+    ( module HHLO.Autograd.Core+    , module HHLO.Autograd.Grad+    , module HHLO.Autograd.Rules+    ) where++import HHLO.Autograd.Core+import HHLO.Autograd.Grad+import HHLO.Autograd.Rules

@@ -0,0 +1,294 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Core+    ( BTensor(..)+    , CotangentMap+    , bconstant+    , badd+    , bsubtract+    , bmultiply+    , bdivide+    , bnegate+    , bexp+    , blog+    , bsqrt+    , bsin+    , bcos+    , btranspose+    , breshape+    , bbroadcastInDim+    , breduceSum+    , bdot+    , babs+    , btanh+    , bmaximum+    , bminimum+    , bselect+    , bslice+    , bpad+    , bconcatenate+    , bconvert+    , bcompareGE+    , btoTyped+    , bfromTyped+    , reifyShape+    , accumulate+    ) where++import Data.Int (Int64)+import Data.Proxy+import qualified Data.Text as T+import GHC.TypeLits+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder++-- ---------------------------------------------------------------------------+-- Backward tensor: runtime-typed tensor handle for the backward pass.+-- ---------------------------------------------------------------------------++-- | A tensor value used during the backward pass, carrying only runtime+-- shape and dtype information.+data BTensor = BTensor+    { btVid  :: !ValueId+    , btType :: !TensorType+    } deriving (Eq, Show)++-- | Map from forward value ID to its accumulated cotangent.+type CotangentMap = Map ValueId BTensor++-- ---------------------------------------------------------------------------+-- Existential shape reification+-- ---------------------------------------------------------------------------++-- | Reify a runtime shape list into a type-level 'KnownShape' constraint.+reifyShape :: [Integer] -> (forall (s :: Shape). KnownShape s => Proxy s -> r) -> r+reifyShape [] k = k (Proxy @'[])+reifyShape (n : ns) k =+    case someNatVal (fromIntegral n) of+        Just (SomeNat (_ :: Proxy n)) ->+            reifyShape ns $ \(_ :: Proxy ns) ->+                k (Proxy @(n ': ns))+        Nothing -> error "autograd-hhlo: negative dimension in reifyShape"++-- ---------------------------------------------------------------------------+-- Conversions between BTensor and typed Tensor+-- ---------------------------------------------------------------------------++-- | Convert a typed 'Tensor' to a runtime-typed 'BTensor'.+bfromTyped :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> BTensor+bfromTyped (Tensor vid) = BTensor vid (tensorType (Proxy @s) (Proxy @d))++-- | Convert a 'BTensor' to a typed 'Tensor', checking that the runtime+-- type matches the expected compile-time type.+btoTyped :: forall s d. (KnownShape s, KnownDType d) => BTensor -> Tensor s d+btoTyped (BTensor vid ttype) =+    let expected = tensorType (Proxy @s) (Proxy @d)+    in if ttype == expected+        then Tensor vid+        else error $ T.unpack $+            "autograd-hhlo: type mismatch in btoTyped. Expected "+            <> T.pack (show expected) <> ", got " <> T.pack (show ttype)++-- ---------------------------------------------------------------------------+-- Cotangent accumulation+-- ---------------------------------------------------------------------------++-- | Add a new cotangent to an existing entry in the map, or insert it+-- if none exists yet.+accumulate :: CotangentMap -> ValueId -> BTensor -> Builder CotangentMap+accumulate cmap vid newBar =+    case Map.lookup vid cmap of+        Just existingBar -> do+            summed <- badd existingBar newBar+            return $! Map.insert vid summed cmap+        Nothing ->+            return $! Map.insert vid newBar cmap++-- ---------------------------------------------------------------------------+-- Primitive backward operations (emit raw StableHLO)+-- ---------------------------------------------------------------------------++bconstant :: TensorType -> Double -> Builder BTensor+bconstant t val = do+    let shp = ttShape t+        dt  = ttDType t+        numElems = product shp+    vid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements shp dt (replicate (fromIntegral numElems) val)] t+    return (BTensor vid t)++badd :: BTensor -> BTensor -> Builder BTensor+badd (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.add" [x, y] [t, t] [] t+    return (BTensor vid t)++bsubtract :: BTensor -> BTensor -> Builder BTensor+bsubtract (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.subtract" [x, y] [t, t] [] t+    return (BTensor vid t)++bmultiply :: BTensor -> BTensor -> Builder BTensor+bmultiply (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.multiply" [x, y] [t, t] [] t+    return (BTensor vid t)++bdivide :: BTensor -> BTensor -> Builder BTensor+bdivide (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.divide" [x, y] [t, t] [] t+    return (BTensor vid t)++bnegate :: BTensor -> Builder BTensor+bnegate (BTensor x t) = do+    vid <- emitOp "stablehlo.negate" [x] [t] [] t+    return (BTensor vid t)++bexp :: BTensor -> Builder BTensor+bexp (BTensor x t) = do+    vid <- emitOp "stablehlo.exponential" [x] [t] [] t+    return (BTensor vid t)++blog :: BTensor -> Builder BTensor+blog (BTensor x t) = do+    vid <- emitOp "stablehlo.log" [x] [t] [] t+    return (BTensor vid t)++bsqrt :: BTensor -> Builder BTensor+bsqrt (BTensor x t) = do+    vid <- emitOp "stablehlo.sqrt" [x] [t] [] t+    return (BTensor vid t)++bsin :: BTensor -> Builder BTensor+bsin (BTensor x t) = do+    vid <- emitOp "stablehlo.sine" [x] [t] [] t+    return (BTensor vid t)++bcos :: BTensor -> Builder BTensor+bcos (BTensor x t) = do+    vid <- emitOp "stablehlo.cosine" [x] [t] [] t+    return (BTensor vid t)++-- | Transpose a BTensor using a permutation.+btranspose :: BTensor -> [Int64] -> TensorType -> Builder BTensor+btranspose (BTensor x inType) perm outType = do+    let permAttr = AttrIntList "permutation" perm+    vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType+    return (BTensor vid outType)++-- | Reshape a BTensor.+breshape :: BTensor -> TensorType -> Builder BTensor+breshape (BTensor x inType) outType = do+    vid <- emitOp "stablehlo.reshape" [x] [inType] [] outType+    return (BTensor vid outType)++-- | Broadcast a BTensor to a new shape with explicit dimension mapping.+bbroadcastInDim :: BTensor -> [Int64] -> TensorType -> Builder BTensor+bbroadcastInDim (BTensor x inType) dims outType = do+    vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType]+        [AttrIntList "dims" (fromIntegral <$> dims)] outType+    return (BTensor vid outType)++-- | Element-wise selection between two BTensors based on a boolean predicate.+bselect :: BTensor -> BTensor -> BTensor -> TensorType -> Builder BTensor+bselect (BTensor p predType) (BTensor t _) (BTensor f _) outType = do+    vid <- emitOp "stablehlo.select" [p, t, f] [predType, outType, outType] [] outType+    return (BTensor vid outType)++-- | Slice a BTensor (forward operation wrapper).+bslice :: BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bslice (BTensor x inType) start limit stride outType = do+    let startAttr = AttrRaw $ "start_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> start) <> ">"+        limitAttr = AttrRaw $ "limit_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> limit) <> ">"+        strideAttr = AttrRaw $ "strides = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> stride) <> ">"+    vid <- emitOp "stablehlo.slice" [x] [inType] [startAttr, limitAttr, strideAttr] outType+    return (BTensor vid outType)++-- | Pad a BTensor with edge and interior padding.+bpad :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bpad (BTensor x inType) (BTensor padVal padType) low high interior outType = do+    let lowAttr  = AttrRaw $ "edge_padding_low = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> low) <> ">"+        highAttr = AttrRaw $ "edge_padding_high = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> high) <> ">"+        intAttr  = AttrRaw $ "interior_padding = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> interior) <> ">"+    vid <- emitOp "stablehlo.pad" [x, padVal] [inType, padType] [lowAttr, highAttr, intAttr] outType+    return (BTensor vid outType)++-- | Concatenate multiple BTensors along a dimension.+bconcatenate :: [BTensor] -> Int64 -> TensorType -> Builder BTensor+bconcatenate inputs dim outType = do+    let vids = map btVid inputs+        types = map btType inputs+        dimAttr = AttrInt "dimension" (fromIntegral dim)+    vid <- emitOp "stablehlo.concatenate" vids types [dimAttr] outType+    return (BTensor vid outType)++-- | Reduce-sum a BTensor over specified dimensions.+breduceSum :: BTensor -> [Int] -> TensorType -> Builder BTensor+breduceSum (BTensor x inType) dims outType = do+    let elemType = TensorType [] (ttDType inType)+    zeroVid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [] (ttDType inType) [0.0]] elemType+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @( 'F32)+        b <- arg @'[] @( 'F32)+        sumVid <- emitOp "stablehlo.add"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [sumVid] [elemType]+    vid <- emitOpRegions "stablehlo.reduce"+            [x, zeroVid]+            [inType, elemType]+            [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)++-- | Matrix multiplication of two BTensors.+bdot :: BTensor -> BTensor -> TensorType -> Builder BTensor+bdot (BTensor x t1) (BTensor y t2) outType = do+    vid <- emitOp "stablehlo.dot" [x, y] [t1, t2] [] outType+    return (BTensor vid outType)++-- | Element-wise absolute value.+babs :: BTensor -> Builder BTensor+babs (BTensor x t) = do+    vid <- emitOp "stablehlo.abs" [x] [t] [] t+    return (BTensor vid t)++-- | Element-wise hyperbolic tangent.+btanh :: BTensor -> Builder BTensor+btanh (BTensor x t) = do+    vid <- emitOp "stablehlo.tanh" [x] [t] [] t+    return (BTensor vid t)++-- | Element-wise maximum.+bmaximum :: BTensor -> BTensor -> Builder BTensor+bmaximum (BTensor x t1) (BTensor y t2) = do+    vid <- emitOp "stablehlo.maximum" [x, y] [t1, t2] [] t1+    return (BTensor vid t1)++-- | Element-wise minimum.+bminimum :: BTensor -> BTensor -> Builder BTensor+bminimum (BTensor x t1) (BTensor y t2) = do+    vid <- emitOp "stablehlo.minimum" [x, y] [t1, t2] [] t1+    return (BTensor vid t1)++-- | Type conversion.+bconvert :: BTensor -> TensorType -> Builder BTensor+bconvert (BTensor x inType) outType = do+    vid <- emitOp "stablehlo.convert" [x] [inType] [] outType+    return (BTensor vid outType)++-- | Greater-than-or-equal comparison (returns boolean tensor).+bcompareGE :: BTensor -> BTensor -> Builder BTensor+bcompareGE (BTensor x t1) (BTensor y t2) = do+    let boolType = TensorType (ttShape t1) Bool+    vid <- emitOp "stablehlo.compare" [x, y] [t1, t2]+        [AttrRaw "comparison_direction = #stablehlo<comparison_direction GE>"] boolType+    return (BTensor vid boolType)

@@ -0,0 +1,195 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Grad+    ( grad+    , gradModule+    , vjp+    , vjpModule+    ) where++import Data.List (foldl')+import Data.Proxy+import Data.Foldable (foldlM)+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)+import HHLO.Core.Types+import HHLO.EDSL.Ops (constant)+import HHLO.IR.AST+import HHLO.IR.Builder++import HHLO.Autograd.Core+import HHLO.Autograd.Rules++-- ---------------------------------------------------------------------------+-- Safe API: gradModule / vjpModule+-- ---------------------------------------------------------------------------++-- | Compute the gradient of a scalar-valued function as a standalone+-- 'Module'.+--+-- The resulting module takes a single tensor argument and returns its+-- gradient (a tensor of the same shape).+gradModule :: forall s d.+              (KnownShape s, KnownDType d)+           => (Tensor s d -> Builder (Tensor '[] d))+           -> Module+gradModule f = Module [gradFunction]+  where+    inType = tensorType (Proxy @s) (Proxy @d)+    arg0   = FuncArg "arg0" inType++    -- Capture the forward trace in isolation.+    forwardFunc :: Function+    forwardFunc = runBuilder @'[] @d "forward" [arg0] $ do+        input <- arg @s @d+        f input++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    -- Build the combined forward+backward function.+    gradFunction :: Function+    gradFunction = runBuilder @s @d "main" [arg0] $ do+        input <- arg @s @d+        y     <- f input++        -- Seed cotangent: 1.0 scalar.+        seed <- constant @'[] @d 1.0+        let initMap = Map.singleton (tensorValue y) (bfromTyped seed)++        -- Backpropagate through forward ops in reverse execution order.+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        -- Extract gradient w.r.t. input.+        case Map.lookup (tensorValue input) finalMap of+            Just bt  -> return (btoTyped @s bt)+            Nothing  -> error "autograd-hhlo: gradient not found for input"++-- | Vector-Jacobian product as a standalone 'Module'.+vjpModule :: forall s t d.+             (KnownShape s, KnownShape t, KnownDType d)+          => (Tensor s d -> Builder (Tensor t d))+          -> Module+vjpModule f = Module [vjpFunction]+  where+    inType  = tensorType (Proxy @s) (Proxy @d)+    seedType = tensorType (Proxy @t) (Proxy @d)+    arg0 = FuncArg "arg0" inType+    arg1 = FuncArg "arg1" seedType++    forwardFunc :: Function+    forwardFunc = runBuilder @t @d "forward" [arg0] $ do+        input <- arg @s @d+        f input++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    vjpFunction :: Function+    vjpFunction = runBuilder @s @d "main" [arg0, arg1] $ do+        input <- arg @s @d+        seed  <- arg @t @d+        _y    <- f input++        let initMap = Map.singleton (tensorValue _y) (bfromTyped seed)+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        case Map.lookup (tensorValue input) finalMap of+            Just bt -> return (btoTyped @s bt)+            Nothing -> error "autograd-hhlo: vjp gradient not found for input"++-- ---------------------------------------------------------------------------+-- In-place API: grad / vjp+-- ---------------------------------------------------------------------------++-- | Reverse-mode gradient combinator that works inside an existing+-- 'Builder' context.+--+-- This builds the gradient module in isolation and inlines its operations+-- into the current builder, remapping value IDs so they remain globally+-- consistent.+--+-- Example:+-- > buildModule @1 @1 "grad_f" $ \x -> grad (\y -> sumAll (multiply y y)) x+grad :: forall s d.+        (KnownShape s, KnownDType d)+     => (Tensor s d -> Builder (Tensor '[] d))+     -> Tensor s d+     -> Builder (Tensor s d)+grad f input = do+    let gradMod = gradModule f+        gradFunc = head (moduleFunctions gradMod)+    gradVid <- inlineFunction gradFunc (Map.singleton (ValueId (-1)) (tensorValue input))+    return (Tensor gradVid)++-- | Vector-Jacobian product inside an existing 'Builder' context.+vjp :: forall s t d.+       (KnownShape s, KnownShape t, KnownDType d)+    => (Tensor s d -> Builder (Tensor t d))+    -> Tensor s d+    -> Tensor t d+    -> Builder (Tensor s d)+vjp f input seed = do+    let vjpMod = vjpModule f+        vjpFunc = head (moduleFunctions vjpMod)+    gradVid <- inlineFunction vjpFunc $ Map.fromList+        [ (ValueId (-1), tensorValue input)+        , (ValueId (-2), tensorValue seed)+        ]+    return (Tensor gradVid)++-- ---------------------------------------------------------------------------+-- Internal helpers+-- ---------------------------------------------------------------------------++-- | Fold a backward step over a list of forward operations, threading the+-- cotangent map through the Builder monad.+foldlBackward :: (Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor))+              -> Map ValueId BTensor+              -> [Operation]+              -> Builder (Map ValueId BTensor)+foldlBackward _ cmap [] = return cmap+foldlBackward step cmap (op:ops) = do+    cmap' <- step cmap op+    foldlBackward step cmap' ops++-- | Inline a function's body operations into the current builder,+-- remapping value IDs so they don't conflict with existing values.+-- The argument mapping maps function argument IDs (negative) to the+-- current builder's actual value IDs.+inlineFunction :: Function -> Map ValueId ValueId -> Builder ValueId+inlineFunction func argMap = do+    finalMap <- foldlM inlineOp argMap (funcBody func)+    let retId = head (funcReturnVids func)+    return $ Map.findWithDefault retId retId finalMap+  where+    inlineOp :: Map ValueId ValueId -> Operation -> Builder (Map ValueId ValueId)+    inlineOp idMap op = do+        let remap vid = Map.findWithDefault vid vid idMap+            newOperands = map remap (opOperands op)+            newRegions = map (remapRegion idMap) (opRegions op)+        newVids <- emitOpRegionsN (opName op) newOperands (opOperandTypes op) (opAttributes op) newRegions (opResultTypes op)+        let updates = zip (opResults op) newVids+        return $ foldl' (\m (old, new) -> Map.insert old new m) idMap updates++    remapRegion :: Map ValueId ValueId -> Region -> Region+    remapRegion idMap (Region blocks) = Region (map (remapBlock idMap) blocks)++    remapBlock :: Map ValueId ValueId -> Block -> Block+    remapBlock idMap (Block blockArgs blockOps) =+        Block blockArgs (map (remapOp idMap) blockOps)++    remapOp :: Map ValueId ValueId -> Operation -> Operation+    remapOp idMap op =+        let remapOpVid vid = Map.findWithDefault vid vid idMap+        in op+            { opOperands = map remapOpVid (opOperands op)+            , opResults  = map remapOpVid (opResults op)+            }++    -- Note: block arguments are local to the block and use negative IDs+    -- allocated by runBlockBuilder. They are self-contained and don't+    -- need remapping because they stay within their block scope.

@@ -0,0 +1,599 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Rules+    ( backwardStep+    ) where++import Data.Int (Int64)+import Data.List (sortOn)+import Data.Maybe (isNothing)+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)++import HHLO.IR.AST+import HHLO.IR.Builder++import HHLO.Autograd.Core++-- ---------------------------------------------------------------------------+-- Main backward step dispatcher+-- ---------------------------------------------------------------------------++-- | Process one forward operation in reverse order, propagating cotangents.+backwardStep :: Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor)+backwardStep cmap op+    | all isNothing resultBars = return cmap+    | otherwise = case opName op of+        "stablehlo.add"          -> vjpAdd op resultBars cmap+        "stablehlo.subtract"     -> vjpSubtract op resultBars cmap+        "stablehlo.multiply"     -> vjpMultiply op resultBars cmap+        "stablehlo.divide"       -> vjpDivide op resultBars cmap+        "stablehlo.negate"       -> vjpNegate op resultBars cmap+        "stablehlo.exponential"  -> vjpExponential op resultBars cmap+        "stablehlo.log"          -> vjpLog op resultBars cmap+        "stablehlo.sqrt"         -> vjpSqrt op resultBars cmap+        "stablehlo.sine"         -> vjpSine op resultBars cmap+        "stablehlo.cosine"       -> vjpCosine op resultBars cmap+        "stablehlo.power"        -> vjpPower op resultBars cmap+        "stablehlo.reshape"      -> vjpReshape op resultBars cmap+        "stablehlo.transpose"    -> vjpTranspose op resultBars cmap+        "stablehlo.broadcast_in_dim" -> vjpBroadcastInDim op resultBars cmap+        "stablehlo.reduce"       -> vjpReduce op resultBars cmap+        "stablehlo.dot"          -> vjpDot op resultBars cmap+        "stablehlo.select"       -> vjpSelect op resultBars cmap+        "stablehlo.slice"        -> vjpSlice op resultBars cmap+        "stablehlo.pad"          -> vjpPad op resultBars cmap+        "stablehlo.concatenate"  -> vjpConcatenate op resultBars cmap+        "stablehlo.gather"       -> vjpGather op resultBars cmap+        "stablehlo.scatter"      -> vjpScatter op resultBars cmap+        "stablehlo.convert"      -> vjpConvert op resultBars cmap+        "stablehlo.abs"          -> vjpAbs op resultBars cmap+        "stablehlo.maximum"      -> vjpMaximum op resultBars cmap+        "stablehlo.minimum"      -> vjpMinimum op resultBars cmap+        "stablehlo.tanh"         -> vjpTanh op resultBars cmap+        "stablehlo.constant"     -> return cmap  -- constants have zero gradient+        "stablehlo.return"       -> return cmap  -- internal terminator+        "stablehlo.compare"      -> return cmap  -- comparisons have zero gradient+        "stablehlo.floor"        -> return cmap  -- non-differentiable+        "stablehlo.ceil"         -> return cmap  -- non-differentiable+        "stablehlo.sort"         -> error "autograd-hhlo: sort/topK is not differentiable"+        _ -> error $ T.unpack $ "autograd-hhlo: no VJP rule for " <> opName op+  where+    resultBars = map (\r -> Map.lookup r cmap) (opResults op)++-- ---------------------------------------------------------------------------+-- VJP helpers+-- ---------------------------------------------------------------------------++getResultBar :: [Maybe BTensor] -> Maybe BTensor+getResultBar [mb] = mb+getResultBar _    = Nothing++operandBT :: Operation -> Int -> BTensor+operandBT op i = BTensor (opOperands op !! i) (opOperandTypes op !! i)++-- ---------------------------------------------------------------------------+-- Arithmetic VJP rules+-- ---------------------------------------------------------------------------++vjpAdd :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpAdd op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            yVid = opOperands op !! 1+        cmap' <- accumulate cmap xVid bar+        accumulate cmap' yVid bar+    Nothing -> return cmap++vjpSubtract :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSubtract op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            yVid = opOperands op !! 1+        negBar <- bnegate bar+        cmap' <- accumulate cmap xVid bar+        accumulate cmap' yVid negBar+    Nothing -> return cmap++vjpMultiply :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMultiply op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        dx <- bmultiply bar y+        dy <- bmultiply bar x+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpDivide :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDivide op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        -- dx = bar / y+        dx <- bdivide bar y+        -- dy = -bar * x / (y * y)+        negBar <- bnegate bar+        t1 <- bmultiply negBar x+        ySq <- bmultiply y y+        dy <- bdivide t1 ySq+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpNegate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpNegate op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        negBar <- bnegate bar+        accumulate cmap (opOperands op !! 0) negBar+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Math VJP rules+-- ---------------------------------------------------------------------------++vjpExponential :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpExponential op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        -- Forward result y = exp(x), so dy/dx = exp(x) = y.+        -- We need the forward result value. For now, recompute exp(x).+        let x = operandBT op 0+        y <- bexp x+        dx <- bmultiply bar y+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpLog :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpLog op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        -- dx = bar / x+        dx <- bdivide bar x+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpSqrt :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSqrt op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        -- dx = bar / (2 * sqrt(x))+        two <- bconstant (btType x) 2.0+        s <- bsqrt x+        t1 <- bmultiply two s+        dx <- bdivide bar t1+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpSine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSine op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        c <- bcos x+        dx <- bmultiply bar c+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpCosine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpCosine op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        s <- bsin x+        negS <- bnegate s+        dx <- bmultiply bar negS+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpPower :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpPower op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        -- For z = x^y:+        -- dz/dx = y * x^(y-1)+        -- dz/dy = x^y * log(x)+        one <- bconstant (btType y) 1.0+        yMinus1 <- bsubtract y one+        xPow <- bpower x yMinus1+        dx <- bmultiply y xPow+        z <- bpower x y+        logX <- blog x+        dy <- bmultiply z logX+        dx' <- bmultiply bar dx+        dy' <- bmultiply bar dy+        cmap' <- accumulate cmap (btVid x) dx'+        accumulate cmap' (btVid y) dy'+    Nothing -> return cmap++-- Helper for power (not exported, used only inside VJP)+bpower :: BTensor -> BTensor -> Builder BTensor+bpower (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.power" [x, y] [t, t] [] t+    return (BTensor vid t)++-- ---------------------------------------------------------------------------+-- Shape manipulation VJP rules+-- ---------------------------------------------------------------------------++vjpReshape :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReshape op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- Gradient is a reshape back to the input shape.+        dx <- breshape bar xType+        accumulate cmap xVid dx+    Nothing -> return cmap++vjpTranspose :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpTranspose op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- Extract permutation from attributes.+        perm <- case findPermutation (opAttributes op) of+            Just p  -> return p+            Nothing -> error "autograd-hhlo: transpose missing permutation attribute"+        -- Inverse permutation.+        let invPerm = invertPerm perm+        dx <- btranspose bar invPerm xType+        accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findPermutation :: [Attribute] -> Maybe [Int64]+    findPermutation [] = Nothing+    findPermutation (AttrIntList "permutation" p : _) = Just p+    findPermutation (_ : attrs) = findPermutation attrs++    invertPerm :: [Int64] -> [Int64]+    invertPerm perm = map fst $ sortOn snd $ zip ([0 ..] :: [Int64]) perm++vjpBroadcastInDim :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpBroadcastInDim op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+            outType = head (opResultTypes op)+        -- Extract broadcast dims.+        dims <- case findDims (opAttributes op) of+            Just d  -> return d+            Nothing -> error "autograd-hhlo: broadcast_in_dim missing dims attribute"+        -- Gradient: reduce over the broadcasted dimensions.+        let outRank = length (ttShape outType) :: Int+            broadcastDims = map fromIntegral dims :: [Int]+            reduceDims = filter (`notElem` broadcastDims) [0 .. outRank - 1]+        if null reduceDims+            then accumulate cmap xVid bar+            else do+                dx <- breduceSum bar reduceDims xType+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findDims [] = Nothing+    findDims (AttrIntList "dims" d : _) = Just d+    findDims (_ : attrs) = findDims attrs++-- ---------------------------------------------------------------------------+-- Reduction VJP rules+-- ---------------------------------------------------------------------------++vjpReduce :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduce op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- We only support sum reductions for now.+        -- Detect sum by looking for stablehlo.add in the reduction region.+        let isSum = any regionHasAdd (opRegions op)+            regionHasAdd (Region blocks) = any blockHasAdd blocks+            blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops+        if not isSum+            then error "autograd-hhlo: only sum reductions are supported"+            else do+                dims <- case findDimensions (opAttributes op) of+                    Just d  -> return $ map fromIntegral d+                    Nothing -> error "autograd-hhlo: reduce missing dimensions attribute"+                -- Gradient: broadcast the cotangent back to the input shape.+                dx <- broadcastLike bar xType dims+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findDimensions [] = Nothing+    findDimensions (AttrRaw raw : _) =+        if "dimensions" `T.isPrefixOf` T.strip raw+        then Just $ parseIntList raw+        else Nothing+    findDimensions (AttrIntList "dimensions" dims : _) = Just dims+    findDimensions (_ : attrs) = findDimensions attrs++    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map (read . T.unpack . T.strip) nums++    -- Broadcast a reduced cotangent back to the original shape.+    broadcastLike :: BTensor -> TensorType -> [Int] -> Builder BTensor+    broadcastLike bar inType dims = do+        let inRank = length (ttShape inType)+            -- Build broadcast dims: map each reduced dim to its position.+            -- For reduce over dims [0,1] of a 3-D tensor, we broadcast+            -- the scalar/shape to the original shape.+            -- We use broadcast_in_dim with dims = remaining_dims.+            remainingDims = filter (`notElem` dims) [0 .. inRank - 1]+        if length remainingDims == inRank+            then return bar  -- No reduction happened+            else do+                bbroadcastInDim bar (map fromIntegral remainingDims) inType++-- ---------------------------------------------------------------------------+-- Linear algebra VJP rules+-- ---------------------------------------------------------------------------++vjpDot :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDot op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+            xType = btType x+            yType = btType y+        -- For C = A @ B:+        -- dA = dC @ B^T+        -- dB = A^T @ dC+        -- We need to transpose the appropriate dimensions.+        -- For simplicity, assume standard 2-D matmul.+        let xShape = ttShape xType+            yShape = ttShape yType+        if length xShape == 2 && length yShape == 2+            then do+                bT <- btranspose y [1, 0] (TensorType (reverse $ ttShape yType) (ttDType yType))+                aT <- btranspose x [1, 0] (TensorType (reverse xShape) (ttDType xType))+                da <- bdot bar bT xType+                db <- bdot aT bar yType+                cmap' <- accumulate cmap (btVid x) da+                accumulate cmap' (btVid y) db+            else+                error "autograd-hhlo: dot VJP only supports 2-D matrices for now"+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Selection VJP rule+-- ---------------------------------------------------------------------------++vjpSelect :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSelect op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let predVid  = opOperands op !! 0+            trueVid  = opOperands op !! 1+            falseVid = opOperands op !! 2+            predType = opOperandTypes op !! 0+            valType  = opOperandTypes op !! 1+        zero <- bconstant valType 0.0+        -- dx = select(pred, bar, 0)+        dx <- bselect (BTensor predVid predType) bar zero valType+        -- dy = select(pred, 0, bar)+        dy <- bselect (BTensor predVid predType) zero bar valType+        cmap' <- accumulate cmap trueVid dx+        accumulate cmap' falseVid dy+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Slice VJP rule+-- ---------------------------------------------------------------------------++vjpSlice :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSlice op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+            (start, limit, stride) = parseSliceAttrs (opAttributes op)+            xShape = ttShape xType+            rank = length xShape+            low = start+            isUnitStride = all (== 1) stride+        if isUnitStride+            then do+                let high' = map fromIntegral (zipWith (-) xShape (map fromIntegral limit :: [Integer])) :: [Int64]+                    interior = replicate rank (0 :: Int64)+                zero <- bconstant xType 0.0+                dx <- bpad bar zero low high' interior xType+                accumulate cmap xVid dx+            else do+                -- General stride: use interior padding = stride - 1+                let start' = map fromIntegral start :: [Integer]+                    limit' = map fromIntegral limit :: [Integer]+                    stride' = map fromIntegral stride :: [Integer]+                    high' = map fromIntegral (zipWith (-) xShape (zipWith (+) start' (zipWith (*) (zipWith (-) limit' start') stride'))) :: [Int64]+                    interior = map (\s -> max 0 (s - 1)) stride :: [Int64]+                zero <- bconstant xType 0.0+                dx <- bpad bar zero low high' interior xType+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    parseSliceAttrs :: [Attribute] -> ([Int64], [Int64], [Int64])+    parseSliceAttrs attrs =+        let s = findAttr "start_indices" attrs+            l = findAttr "limit_indices" attrs+            st = findAttr "strides" attrs+        in (s, l, st)++    findAttr :: Text -> [Attribute] -> [Int64]+    findAttr _ [] = error "autograd-hhlo: slice missing attribute"+    findAttr name (AttrRaw raw : rest) =+        let key = name <> " = array<i64:"+        in if key `T.isInfixOf` raw+            then parseIntList raw+            else findAttr name rest+    findAttr name (_ : rest) = findAttr name rest++    parseIntList :: Text -> [Int64]+    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Concatenate VJP rule+-- ---------------------------------------------------------------------------++vjpConcatenate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConcatenate op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let inputVids = opOperands op+            inputTypes = opOperandTypes op+        dim <- case findConcatDim (opAttributes op) of+            Just d  -> return d+            Nothing -> error "autograd-hhlo: concatenate missing dimension attribute"+        -- Split bar along concat dimension into slices matching each input+        let inputSizes = map (!! dim) (map ttShape inputTypes)+        splitAndAccumulate bar dim 0 inputVids inputTypes inputSizes cmap+    Nothing -> return cmap+  where+    findConcatDim :: [Attribute] -> Maybe Int+    findConcatDim [] = Nothing+    findConcatDim (AttrInt "dimension" d : _) = Just (fromIntegral d)+    findConcatDim (_ : rest) = findConcatDim rest++    splitAndAccumulate :: BTensor -> Int -> Integer -> [ValueId] -> [TensorType] -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+    splitAndAccumulate _ _ _ [] [] [] acc = return acc+    splitAndAccumulate bar dim offset (vid:vids) (itype:itypes) (sz:szs) acc = do+        let shape = ttShape itype+            start = replicate (length shape) (0 :: Integer)+            limit = zipWith (\i s -> if i == dim then offset + s else s) [0..] shape+            stride = replicate (length shape) (1 :: Integer)+        piece <- bslice bar (map fromIntegral start) (map fromIntegral limit) (map fromIntegral stride) itype+        acc' <- accumulate acc vid piece+        splitAndAccumulate bar dim (offset + sz) vids itypes szs acc'+    splitAndAccumulate _ _ _ _ _ _ _ = error "autograd-hhlo: concatenate operand mismatch"++-- ---------------------------------------------------------------------------+-- Pad VJP rule+-- ---------------------------------------------------------------------------++vjpPad :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpPad op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        (low, _high, interior) <- parsePadAttrs (opAttributes op)+        let xShape = ttShape xType+            stride = map (+ 1) interior :: [Int64]+            limit = zipWith3 (\l s sz -> l + sz * s) low stride (map fromIntegral xShape) :: [Int64]+        dx <- bslice bar low limit stride xType+        accumulate cmap xVid dx+    Nothing -> return cmap+  where+    parsePadAttrs :: [Attribute] -> Builder ([Int64], [Int64], [Int64])+    parsePadAttrs attrs =+        case (findAttr "edge_padding_low" attrs,+              findAttr "edge_padding_high" attrs,+              findAttr "interior_padding" attrs) of+            (Just l, Just h, Just i) -> return (l, h, i)+            _ -> error "autograd-hhlo: pad missing padding attributes"++    findAttr :: Text -> [Attribute] -> Maybe [Int64]+    findAttr _ [] = Nothing+    findAttr name (AttrRaw raw : rest) =+        let key = name <> " = array<i64:"+        in if key `T.isInfixOf` raw+            then Just (parseIntList raw)+            else findAttr name rest+    findAttr name (_ : rest) = findAttr name rest++    parseIntList :: Text -> [Int64]+    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Gather / Scatter VJP rules (stubs)+-- ---------------------------------------------------------------------------++vjpGather :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpGather _ _ _ = error "autograd-hhlo: gather VJP not yet implemented"++vjpScatter :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpScatter _ _ _ = error "autograd-hhlo: scatter VJP not yet implemented"++-- ---------------------------------------------------------------------------+-- Convert VJP rule+-- ---------------------------------------------------------------------------++vjpConvert :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConvert op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        dx <- bconvert bar xType+        accumulate cmap xVid dx+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Abs VJP rule+-- ---------------------------------------------------------------------------++vjpAbs :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpAbs op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        absX <- babs x+        dx <- bdivide x absX+        dx' <- bmultiply bar dx+        accumulate cmap (btVid x) dx'+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Maximum / Minimum VJP rules+-- ---------------------------------------------------------------------------++vjpMaximum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMaximum op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        predGE <- bcompareGE x y+        zero <- bconstant (btType x) 0.0+        dx <- bselect predGE bar zero (btType x)+        dy <- bselect predGE zero bar (btType x)+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpMinimum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMinimum op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        predLE <- bcompareGE y x+        zero <- bconstant (btType x) 0.0+        dx <- bselect predLE bar zero (btType x)+        dy <- bselect predLE zero bar (btType x)+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Tanh VJP rule+-- ---------------------------------------------------------------------------++vjpTanh :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpTanh op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        tanhX <- btanh x+        tanhSq <- bmultiply tanhX tanhX+        one <- bconstant (btType x) 1.0+        oneMinus <- bsubtract one tanhSq+        dx <- bmultiply bar oneMinus+        accumulate cmap (btVid x) dx+    Nothing -> return cmap

@@ -11,6 +11,7 @@     , divide     , matmul     , dotGeneral+    , einsum     , linear     , linearBatched     -- * Unary element-wise ops@@ -40,9 +41,13 @@     , concatenate     , concatenate2     , iota+    , split+    , stack     -- * Reductions     , reduceSum     , reduceSumDim+    , productAll+    , productDim     , reduceWindow     , maxPool     , avgPool@@ -94,6 +99,7 @@     , dynamicSlice     , sort     , convert+    , topK     -- * Selection     , select     -- * Map@@ -131,7 +137,10 @@  import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh, sqrt, sin, cos, tan, floor, ceiling) +import Control.Monad (when) import Data.Int (Int64)+import Data.List (elemIndex)+import Data.Maybe (fromJust) import Data.Proxy import Data.Text (Text) import qualified Data.Text as T@@ -1879,3 +1888,200 @@     y1 <- reshape @'[] @'[1] y     z1 <- reshape @'[] @'[1] z     concatenate @'[1] @'[3] @d 0 [x1, y1, z1]++-- ---------------------------------------------------------------------------+-- Product reductions+-- ---------------------------------------------------------------------------++-- | Product of all elements (reduce over all dimensions).+productAll :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d)+productAll (Tensor x) = do+    let dims = [0 .. fromIntegral (length (shapeVal (Proxy @s))) - 1]+    productDim @s @'[] dims (Tensor x)++-- | Product elements over specific dimensions.+productDim :: forall sFrom sTo d.+              (KnownShape sFrom, KnownShape sTo, KnownDType d)+           => [Int] -> Tensor sFrom d -> Builder (Tensor sTo d)+productDim dims (Tensor x) = do+    let inType   = tensorType (Proxy @sFrom) (Proxy @d)+        outType  = tensorType (Proxy @sTo)   (Proxy @d)+        elemType = tensorType (Proxy @'[])   (Proxy @d)+    -- Init value for stablehlo.reduce (must be scalar)+    oneVid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [] (dtypeVal (Proxy @d)) [1.0]] elemType+    -- Build reduction region: two scalar args, apply stablehlo.multiply+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @d+        b <- arg @'[] @d+        prodVid <- emitOp "stablehlo.multiply"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [prodVid] [elemType]+    vid <- emitOpRegions "stablehlo.reduce"+            [x, oneVid]+            [inType, elemType]+            [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+            [Region [redBlock]]+            outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Split and stack+-- ---------------------------------------------------------------------------++-- | Split a tensor into @n@ equal parts along dimension @dim@.+-- The size of @dim@ in the input must be evenly divisible by @n@.+split :: forall sIn sOut d.+         (KnownShape sIn, KnownShape sOut, KnownDType d)+      => Int64            -- ^ dimension to split along+      -> Int64            -- ^ number of equal splits+      -> Tensor sIn d+      -> Builder [Tensor sOut d]+split dim n t = do+    let sInShape = shapeVal (Proxy @sIn)+        rank = length sInShape+        dimSize = fromIntegral (sInShape !! fromIntegral dim) :: Int+        chunkSize = dimSize `div` fromIntegral n+        stride = replicate rank 1+    when (dimSize `mod` fromIntegral n /= 0) $+        error "split: dimension size not evenly divisible by number of splits"+    Prelude.mapM (\i -> do+        let start = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]]+            limit = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]]+        slice @sIn @sOut @d t start limit stride+        ) [0 .. fromIntegral n - 1]++-- | Stack a list of tensors along a new axis.+-- All inputs must have the same shape. The new axis is inserted at position+-- @dim@ in the output shape.+stack :: forall sIn sOut d.+         (KnownShape sIn, KnownShape sOut, KnownDType d)+      => Int64            -- ^ axis to insert and stack along+      -> [Tensor sIn d]   -- ^ tensors to stack (must be non-empty)+      -> Builder (Tensor sOut d)+stack dim inputs = do+    let n = length inputs+        inShape  = fmap fromIntegral (shapeVal (Proxy @sIn)) :: [Int64]+        outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64]+        dt = dtypeVal (Proxy @d)+        rank = length inShape+        expectedOutShape = take (fromIntegral dim) inShape ++ [fromIntegral n] ++ drop (fromIntegral dim) inShape+        reshapedShape = take (fromIntegral dim) inShape ++ [1] ++ drop (fromIntegral dim) inShape+        inType = tensorType (Proxy @sIn) (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        reshapedType = TensorType (fmap fromIntegral reshapedShape) dt+    when (n == 0) $ error "stack: empty input list"+    when (outShape /= expectedOutShape) $+        error $ "stack: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape+    reshapedVids <- Prelude.mapM (\(Tensor vid) ->+        emitOp "stablehlo.reshape" [vid] [inType] [] reshapedType+        ) inputs+    let dimAttr = AttrInt "dimension" (fromIntegral dim)+    vid <- emitOp "stablehlo.concatenate" reshapedVids (replicate n reshapedType) [dimAttr] outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Top-K+-- ---------------------------------------------------------------------------++-- | Return the top-K values along a dimension, in descending order.+--+-- Note: A @topKWithIndices@ variant is future work; it requires multi-operand+-- sort support (sorting values and their index tensors together).+topK :: forall s sOut d.+        (KnownShape s, KnownShape sOut, KnownDType d)+     => Int64            -- ^ K+     -> Int64            -- ^ dimension to sort along+     -> Tensor s d+     -> Builder (Tensor sOut d)+topK k dim t = do+    let sShape = fmap fromIntegral (shapeVal (Proxy @s)) :: [Int64]+        outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64]+        rank = length sShape+        expectedOutShape = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]]+    when (outShape /= expectedOutShape) $+        error $ "topK: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape+    -- Sort descending+    sorted <- sort @s @d t dim False $ \a b -> greaterThan a b+    -- Slice the first K elements along dim+    let start = replicate rank 0+        limit = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]]+        stride = replicate rank 1+    slice @s @sOut @d sorted start limit stride++-- ---------------------------------------------------------------------------+-- Einsum+-- ---------------------------------------------------------------------------++-- | Einstein summation for two tensors.+--+-- Example: @einsum "ijk,jkl->il" a b@ contracts the @j@ and @k@ dimensions,+-- leaving @i@ from the left operand and @l@ from the right.+--+-- The caller must supply the output shape @sOut@ via type application or+-- inference context. Runtime validation ensures the subscript string agrees+-- with the type-level shapes.+einsum :: forall s1 s2 sOut d.+          (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d)+       => String              -- ^ subscript string, e.g. "ijk,jkl->il"+       -> Tensor s1 d+       -> Tensor s2 d+       -> Builder (Tensor sOut d)+einsum spec (Tensor x) (Tensor y) = do+    let (left, right, out) = parseEinsum spec+        s1Shape = shapeVal (Proxy @s1)+        s2Shape = shapeVal (Proxy @s2)+        sOutShape = shapeVal (Proxy @sOut)+        dt = dtypeVal (Proxy @d)+        rank1 = length s1Shape+        rank2 = length s2Shape+        rankOut = length sOutShape+    when (length left /= rank1) $ error "einsum: left subscript rank mismatch"+    when (length right /= rank2) $ error "einsum: right subscript rank mismatch"+    when (length out /= rankOut) $ error "einsum: output subscript rank mismatch"+    let batchLabels     = [c | c <- left, c `elem` right, c `elem` out]+        contractLabels  = [c | c <- left, c `elem` right, c `notElem` out]+        leftFree        = [c | c <- left, c `elem` out, c `notElem` right]+        rightFree       = [c | c <- right, c `elem` out, c `notElem` left]+        natural         = batchLabels ++ leftFree ++ rightFree+        lhsBatch        = fmap (fromIntegral . fromJust . (`elemIndex` left)) batchLabels+        rhsBatch        = fmap (fromIntegral . fromJust . (`elemIndex` right)) batchLabels+        lhsContract     = fmap (fromIntegral . fromJust . (`elemIndex` left)) contractLabels+        rhsContract     = fmap (fromIntegral . fromJust . (`elemIndex` right)) contractLabels+        -- Build natural shape from actual dimensions+        lookupDim label+            | label `elem` left  = s1Shape !! fromJust (label `elemIndex` left)+            | otherwise          = s2Shape !! fromJust (label `elemIndex` right)+        naturalShape    = fmap (fromIntegral . lookupDim) natural+        naturalType     = TensorType (fmap fromIntegral naturalShape) dt+        inType1         = tensorType (Proxy @s1) (Proxy @d)+        inType2         = tensorType (Proxy @s2) (Proxy @d)+        outType         = tensorType (Proxy @sOut) (Proxy @d)+        batchAttr       = AttrString "batching_dims"+                            ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsBatch) <> "] x ["+                               <> T.intercalate ", " (fmap (T.pack . show) rhsBatch) <> "]")+        contractingAttr = AttrString "contracting_dims"+                            ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsContract) <> "] x ["+                               <> T.intercalate ", " (fmap (T.pack . show) rhsContract) <> "]")+        -- Validate output shape+        expectedOutShape = fmap (fromIntegral . lookupDim) out+    when (fmap fromIntegral sOutShape /= expectedOutShape) $+        error $ "einsum: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show (fmap fromIntegral sOutShape)+    vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]+            [batchAttr, contractingAttr] naturalType+    if natural == out+    then return (Tensor vid)+    else do+        let perm = fmap (fromIntegral . fromJust . (`elemIndex` natural)) out+            permAttr = AttrIntList "permutation" perm+        vidT <- emitOp "stablehlo.transpose" [vid] [naturalType] [permAttr] outType+        return (Tensor vidT)+  where+    parseEinsum s =+        case break (== '-') s of+            (lhsRhs, '-':'>':outRest) ->+                case break (== ',') lhsRhs of+                    (left, ',':right) -> (left, right, outRest)+                    _ -> error "einsum: expected two operands separated by comma"+            _ -> error "einsum: expected -> in subscript string"

@@ -176,10 +176,10 @@         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : () -> " <> prettyResults resultTypes     pretty (Operation "stablehlo.sort" operands operandTypes attrs regions results resultTypes) =-        -- Generic form (has regions; fallback would already use generic, but explicit is clearer).+        -- Generic form: regions must be wrapped in ( ) for PJRT parser compatibility.         prettyResultVids results <> " = \"stablehlo.sort\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"-        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null regions then mempty else " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")")         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) =

@@ -5,7 +5,8 @@ import qualified Test.IR.Pretty as Pretty import qualified Test.IR.Builder as Builder import qualified Test.EDSL.Ops as EDSLOps-import qualified Test.ModuleBuilder as ModuleBuilder+import qualified Test.Autograd.Grad as AutogradGrad+import qualified Test.Autograd.Rules as AutogradRules import qualified Test.Runtime.EndToEnd as EndToEnd import qualified Test.Runtime.EndToEndArithmetic as Arith import qualified Test.Runtime.EndToEndShape as Shape@@ -15,6 +16,7 @@ import qualified Test.Runtime.EndToEndDataMovement as DataMovement import qualified Test.Runtime.EndToEndMultiValue as MultiValue import qualified Test.Runtime.EndToEndSession as Session+import qualified Test.Runtime.EndToEndAutograd as Autograd import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors@@ -38,6 +40,8 @@         [ Pretty.tests         , Builder.tests         , EDSLOps.tests+        , AutogradGrad.tests+        , AutogradRules.tests         , EndToEnd.tests         , Arith.tests         , Shape.tests@@ -47,6 +51,7 @@         , DataMovement.tests         , MultiValue.tests         , Session.tests+        , Autograd.tests         , Buffer.tests         , Async.tests         , Errors.tests

@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Autograd.Grad (tests) where++import Prelude hiding (sqrt)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Pretty+import HHLO.Autograd++import qualified Data.Text as T++tests :: TestTree+tests = testGroup "Autograd.Grad"+    [ testCase "sumOfSquaresGrad" $ do+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.multiply" ("stablehlo.multiply" `T.isInfixOf` text)+        assertBool "should contain stablehlo.add" ("stablehlo.add" `T.isInfixOf` text)+    , testCase "sumOfDoublesGrad" $ do+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.constant" ("stablehlo.constant" `T.isInfixOf` text)+    , testCase "multiplyChainGrad" $ do+        let f x = do+                sq <- multiply x x+                cb <- multiply sq x+                sumAll cb+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should be non-empty" (not $ T.null text)+    , testCase "mixedOpsGrad" $ do+        let f x = do+                sq <- multiply x x+                s <- sqrt sq+                sumAll s+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.sqrt" ("stablehlo.sqrt" `T.isInfixOf` text)+    ]

@@ -0,0 +1,49 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Autograd.Rules (tests) where++import Prelude hiding (negate)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Pretty+import HHLO.Autograd++import qualified Data.Text as T++tests :: TestTree+tests = testGroup "Autograd.Rules"+    [ testCase "vjpAdd" $ do+        let f x = do+                y <- add x x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "contains func.func" ("func.func" `T.isInfixOf` text)+        assertBool "contains return" ("return" `T.isInfixOf` text)+    , testCase "vjpMultiply" $ do+        let f x = do+                y <- multiply x x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpNegate" $ do+        let f x = do+                y <- negate x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpExponential" $ do+        let f x = do+                y <- exponential x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    ]

@@ -610,5 +610,76 @@                     [ FuncArg "arg0" (TensorType [3] F32) ]                     $ do x <- arg @'[3] @'F32; y <- slice1 x 1; return y             assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu+        , testCase "productAll" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                    $ do x <- arg @'[2, 3] @'F32; y <- productAll x; return y+            let rendered = render modu+            assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+            assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered+        , testCase "productDim" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                    $ do x <- arg @'[2, 3] @'F32; y <- productDim @'[2, 3] @'[2] [1] x; return y+            let rendered = render modu+            assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+            assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered+        , testCase "split" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        ys <- split @'[4] @'[2] 0 2 x+                        case ys of+                            (y1:_) -> return y1+                            _ -> error "expected at least one split"+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu+        , testCase "stack" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- stack @'[2] @'[2, 2] 0 [x, y]+                        return z+            let rendered = render modu+            assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered+        , testCase "topK" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- topK @'[4] @'[2] 2 0 x+                        return y+            let rendered = render modu+            assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered+        , testCase "einsum matmul" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32)+                    , FuncArg "arg1" (TensorType [3, 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- einsum "ij,jk->ik" x y+                        return z+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` render modu+        , testCase "einsum transpose output" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32)+                    , FuncArg "arg1" (TensorType [3, 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- einsum "ij,jk->ki" x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+            assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered         ]     ]

@@ -0,0 +1,74 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndAutograd where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Pretty+import HHLO.Autograd+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.Autograd"+    [ testCase "grad sum of squares" $ withPJRTCPU $ \api client -> do+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        -- grad = 2 * x = [2, 4, 6]+        let expected = V.fromList [2.0, 4.0, 6.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of doubles" $ withPJRTCPU $ \api client -> do+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        -- grad = [2, 2, 2]+        let expected = V.fromList [2.0, 2.0, 2.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of exponentials" $ withPJRTCPU $ \api client -> do+        let f x = do e <- exponential x; sumAll e+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        -- grad = exp(x) = [1, 1, 1]+        let expected = V.fromList [1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad matmul" $ withPJRTCPU $ \api client -> do+        let f x = do+                w <- constant @'[3, 2] @'F32 0.5+                y <- matmul x w+                sumAll y+            gradModu = gradModule @'[2, 3] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 6+        -- grad = sum over cols of W = [0.5+0.5, 0.5+0.5, 0.5+0.5] for each row+        -- = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    ]

@@ -257,4 +257,18 @@         [bufOut] <- execute api exec [bufA]         result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)         result @?= V.fromList [0, 1, 0]+    , testCase "topK" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    y <- topK @'[4] @'[2] 2 0 x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [3.0, 1.0, 4.0, 1.0]+        bufIn <- toDeviceF32 api client inp [4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- Descending sort: [4.0, 3.0, 1.0, 1.0], top 2: [4.0, 3.0]+        result @?= V.fromList [4.0, 3.0]     ]

@@ -96,4 +96,45 @@         let expected = V.fromList [3.0, 3.0, 7.5, 7.5]         assertBool "dotGeneral close" $             all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum matmul" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ik" x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 3]+        bufB <- toDeviceF32 api client b [3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "einsum matmul close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum transpose output" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ki" x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 3]+        bufB <- toDeviceF32 api client b [3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        -- Same as matmul but transposed: [[22,49],[28,64]] flattened row-major = [22,49,28,64]+        let expected = V.fromList [22.0, 49.0, 28.0, 64.0]+        assertBool "einsum transpose close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))     ]

@@ -66,4 +66,31 @@         let expected = V.fromList [3.5, 5.5, 11.5, 13.5]         assertBool "avgPool close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "productAll" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productAll x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [720.0]+    , testCase "productDim" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productDim @'[2, 3] @'[2] [1] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- Row products: 1*2*3=6, 4*5*6=120+        result @?= V.fromList [6.0, 120.0]     ]

@@ -93,4 +93,37 @@         [bufOut] <- execute api exec []         result <- fromDeviceF32 api bufOut 4         result @?= V.fromList [0.0, 1.0, 2.0, 3.0]+    , testCase "split" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    ys <- split @'[4] @'[2] 0 2 x+                    case ys of+                        (y1:_) -> return y1+                        _ -> error "expected at least one split"+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "stack" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    z <- stack @'[2] @'[2, 2] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+        bufA <- toDeviceF32 api client a [2]+        bufB <- toDeviceF32 api client b [2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]     ]