hhlo 0.5.0.0 → 0.6.0.0

raw patch · 10 files changed

+626/−11 lines, 10 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +16/−1
• doc/implementation-design.md +23/−6
• doc/progress-and-remaining-work.md +3/−1
• doc/test-suite-documentation.md +25/−2
• hhlo.cabal +1/−1
• src/HHLO/Autograd/Core.hs +103/−0
• src/HHLO/Autograd/Rules.hs +354/−0
• src/HHLO/IR/Pretty.hs +11/−0
• test/Test/Autograd/Rules.hs +32/−0
• test/Test/Runtime/EndToEndAutograd.hs +58/−0

@@ -67,7 +67,6 @@ * 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.@@ -87,3 +86,19 @@ * 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.++## 0.6.0.0 -- 2026-04-28++* **Convolution & pooling VJP rules** — autograd now supports backprop through+  `conv2d`, `transposeConvolution`, `maxPool`, and `avgPool`.+  * `vjpConvolution` / `vjpTransposeConvolution` emit backward input via+    flipped-kernel transposed conv and skip backward-kernel computation when+    the kernel is a constant (the common `gradModule` case).+  * `vjpReduceWindow` supports both sum-based (avgPool) and select-mask-based+    (maxPool) backward passes.+  * New primitive emitters: `bconvolution`, `breverse`.+  * PJRT parser compatibility: `stablehlo.reverse` custom pretty-printer and+    `batch_group_count` / `feature_group_count` attributes on backward convs.+* New E2E autograd tests: `grad conv2d`, `grad maxPool`, `grad avgPool`.+* New unit tests: `vjpConvolution`, `vjpTransposeConvolution`, `vjpReduceWindow`.+* Test count: 187 CPU tests + 6 GPU integration tests.

@@ -444,16 +444,31 @@  No changes to Layers 1–3 are needed. Only Layer 4 needs device enumeration APIs. -### 10.2 Automatic Differentiation+### 10.2 Automatic Differentiation ✅ Implemented -Reverse-mode autodiff can be added as a source-to-source transformation on the `Builder` monad:+Reverse-mode autodiff is implemented as a source-to-source transformation on the `Builder` monad: -1. Record the computation graph during `Builder` execution-2. Traverse the graph backward, emitting gradient ops-3. Use `stablehlo.custom_call` for ops without native gradient definitions+1. **`Builder` graph recording** — `runBuilderWithTrace` records every emitted operation in a `Trace` (a list of `Operation` values).+2. **Backward traversal** — `gradModule` reverses the trace and runs a backward pass: for each forward op, its VJP rule emits gradient ops that propagate cotangents.+3. **VJP rules** — Each primitive op has a rule in `HHLO.Autograd.Rules` that computes how gradients flow through it. Rules exist for 25+ ops including element-wise ops, reductions, matmul, convolution, transpose convolution, and reduce_window (max/avg pool).+4. **Public API** — `grad :: (Tensor s d -> Builder (Tensor '[] d)) -> (Tensor s d -> Builder (Tensor s d))` gives the gradient of a scalar-valued function w.r.t. its input. -This is analogous to JAX's `jax.grad` but operates on the AST level rather than tracing Python.+**Multi-parameter gradient workaround:**+`gradModule` differentiates w.r.t. a single input tensor. To differentiate w.r.t. multiple weight tensors (e.g., for training), pack all parameters into a single tensor via `concatenate`, pass that as the input, and slice it apart inside the builder: +```haskell+trainStep paramsBatch input = do+    let (w1, w2, b1, b2) = unpackParams paramsBatch+    y <- model w1 w2 b1 b2 input+    loss <- mseLoss y target+    return loss+  where+    unpackParams p = (slice @0 @0 @n1 p, slice @0 @n1 @n2 p,+                      slice @0 @(n1+n2) @(n1+n2+m1) p, ...)+```++This pattern is demonstrated in `examples/35-autograd-linear.hs`.+ ### 10.3 Template Haskell Staging  For compile-time MLIR generation (AOT compilation):@@ -501,6 +516,8 @@ | `HHLO.Runtime.Compile` | Compilation | `compile` | | `HHLO.Runtime.Execute` | Sync execution | `execute` | | `HHLO.Runtime.Async` | Async execution | `executeAsync`, `bufferReady`, `awaitBuffers` |+| `HHLO.Autograd.Core` | AD infrastructure | `gradModule`, `grad`, `BTensor`, `badd`, `bconvolution`, `breverse` |+| `HHLO.Autograd.Rules` | VJP rule dispatch | `vjpAdd`, `vjpMatmul`, `vjpConvolution`, `vjpReduceWindow`, ... | | `HHLO.Runtime.Buffer` | Buffer transfer | `toDeviceF32`, `fromDeviceF32`, `bufferDimensions` |  ---

@@ -123,6 +123,7 @@ | **18** | **Backend-agnostic plugin loading** | ✅ **`withPJRT` abstracts CPU/CUDA plugin selection.** | | **19** | **GPU examples & benchmarks** | ✅ **`example-gpu-add` and `example-gpu-matmul-bench` operational.** | | **20** | **Multi-GPU inference scaling** | ✅ **`executeReplicas` runs concurrent `executeOn` across N GPUs. `compileWithOptions` supports `num_replicas`. `example-multi-gpu-inference` verified on 8× RTX 5090.** |+| **21** | **Autograd (reverse-mode AD)** | ✅ **`gradModule`, `grad`, and VJP rules for 25+ ops including `convolution`, `transpose_convolution`, `reduce_window`. 7 E2E autograd tests pass.** |  --- @@ -131,7 +132,7 @@ ### 1. Single-Device Execution (GPU works, multi-GPU inference works) - `executeOn` targets exactly one device. ✅ - `executeReplicas` distributes independent forward passes across multiple GPUs concurrently. ✅-- **Clarification:** HHLO is an **inference-only** framework. We do not have automatic differentiation, gradients, or backpropagation. Multi-GPU means inference scaling only.+- **Clarification:** HHLO now supports **reverse-mode automatic differentiation** via `grad` / `gradModule`. Multi-GPU still means inference scaling only — autograd runs on a single device. See the autograd examples (34–36) and the autograd section in `doc/implementation-design.md`.  ### 2. PJRT CPU v1.16.0 Parser Limitations The specific `libpjrt_cpu.so` build from `zml/pjrt-artifacts` (StableHLO v1.16.0) has a text parser with known gaps:@@ -190,6 +191,7 @@ 13. ~~UNet inference example~~ ✅ Done. 14. ~~Comprehensive test suite~~ ✅ Done. 15. ~~Single-GPU CUDA support~~ ✅ Done.+16. ~~Reverse-mode automatic differentiation~~ ✅ Done.  --- 

@@ -1,7 +1,7 @@ # HHLO Test Suite — Comprehensive Documentation  **Date:** 2026-04-20  -**Test Count:** 115 tests across 13 modules  +**Test Count:** 187 tests across 15 modules   **Framework:** `tasty` + `tasty-hunit`   **Entry Point:** `test/Main.hs` → `test-suite hhlo-test` in `hhlo.cabal` @@ -15,6 +15,7 @@ |------|------|-------|----------------| | **Tier 1 — Golden** | Rendered MLIR text matches expected StableHLO | `Test.IR.Pretty*`, `Test.EDSL.Ops`, `Test.IR.Builder` | ❌ No | | **Tier 2 — E2E Numerical** | Build → Compile → Execute → Verify on CPU | `Test.Runtime.EndToEnd*` | ✅ Yes |+| **Tier 2 — Autograd** | Gradient computation via reverse-mode AD | `Test.Autograd.Rules`, `Test.Runtime.EndToEndAutograd` | ✅ Yes | | **Tier 3 — Runtime Integration** | Buffer metadata, async execution, error handling | `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors` | ✅ Yes |  All E2E tests use the PJRT CPU plugin at `deps/pjrt/libpjrt_cpu.so` (downloaded via `./pjrt_script.sh`).@@ -117,8 +118,18 @@ | `module has func.func wrapper` | `module { func.func @main(...) }` | | `single result type in signature` | `-> tensor<3x4xf32>` | -### 3.6 `test/Test/EDSL/Ops.hs`+### 3.6 `test/Test/Autograd/Rules.hs` +Golden tests for VJP (vector-Jacobian product) rules. Each test builds a small forward computation, applies the VJP rule, and verifies the backward MLIR contains the expected ops.++| Test | Forward Op | Backward Ops Verified |+|------|-----------|----------------------|+| `vjpReduceWindow (avgPool)` | `avgPool` | `broadcast_in_dim`, `divide`, `pad` |+| `vjpConvolution` | `conv2d` | `stablehlo.reverse`, `stablehlo.convolution` |+| `vjpTransposeConvolution` | `transposeConvolution` | `stablehlo.convolution` |++### 3.7 `test/Test/EDSL/Ops.hs`+ The largest golden test module. It exercises virtually every EDSL op and verifies that the rendered MLIR contains the expected StableHLO op name or structural pattern.  **Categories:**@@ -216,6 +227,18 @@ | `transpose` | `transpose` | `[[1,2],[3,4]]` → `[[1,3],[2,4]]` | | `concatenate` | `concatenate` | `[1,2] + [3,4]` → `[1,2,3,4]` | | `iota 1D` | `iota` + `convert` | `[0,1,2,3]` (iota returns `I64`, converted to `F32`) |++### 4.7 `test/Test/Runtime/EndToEndAutograd.hs`++| Test | What it computes | Status |+|------|-----------------|--------|+| `grad sum of squares` | `grad (\x -> sumAll (x * x))` | ✅ Pass |+| `grad sum of doubles` | `grad (\x -> sumAll (x + x))` | ✅ Pass |+| `grad sum of exponentials` | `grad (\x -> sumAll (exp x))` | ✅ Pass |+| `grad matmul` | `grad (\x -> sumAll (matmul x w))` | ✅ Pass |+| `grad avgPool` | `grad (\x -> sumAll (avgPool x))` | ✅ Pass |+| `grad conv2d` | `grad (\x -> sumAll (conv2d x k))` | ✅ Pass |+| `grad maxPool` | `grad (\x -> sumAll (maxPool x))` | ✅ Pass |  --- 

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.5.0.0+version:            0.6.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

@@ -33,14 +33,20 @@     , bconcatenate     , bconvert     , bcompareGE+    , bcompareEQ     , btoTyped     , bfromTyped     , reifyShape     , accumulate+    , breverse+    , bconvolution+    , breduceWindowAdd+    , breduceWindowMax     ) where  import Data.Int (Int64) import Data.Proxy+import Data.Text (Text) import qualified Data.Text as T import GHC.TypeLits import qualified Data.Map.Strict as Map@@ -292,3 +298,100 @@     vid <- emitOp "stablehlo.compare" [x, y] [t1, t2]         [AttrRaw "comparison_direction = #stablehlo<comparison_direction GE>"] boolType     return (BTensor vid boolType)++-- | Equal comparison (returns boolean tensor).+bcompareEQ :: BTensor -> BTensor -> Builder BTensor+bcompareEQ (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 EQ>"] boolType+    return (BTensor vid boolType)++-- | Reverse a tensor along specified dimensions.+breverse :: BTensor -> [Int64] -> TensorType -> Builder BTensor+breverse (BTensor x t) dims outType = do+    let dimsAttr = AttrIntList "dimensions" dims+    vid <- emitOp "stablehlo.reverse" [x] [t] [dimsAttr] outType+    return (BTensor vid outType)++-- | Generic convolution emitter.+--+-- This is the low-level primitive used by VJP rules.  It emits a+-- @stablehlo.convolution@ with fully-specified dimension numbers and+-- window attributes.+bconvolution :: BTensor -> BTensor -> Text -> Text -> [Attribute] -> TensorType -> Builder BTensor+bconvolution (BTensor lhs lhsType) (BTensor rhs rhsType) dimNums windowStr extraAttrs outType = do+    let attrs =+            [ AttrString "dim_numbers" dimNums+            , AttrString "window" windowStr+            ] ++ extraAttrs+    vid <- emitOp "stablehlo.convolution"+            [lhs, rhs]+            [lhsType, rhsType]+            attrs+            outType+    return (BTensor vid outType)++-- | Reduce a BTensor over specified window dimensions with @add@.+breduceWindowAdd :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor+breduceWindowAdd (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do+    let elemType = TensorType [] (ttDType inType)+    -- Build the add reduction region.+    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]++    let windowAttr  = AttrRaw $ "window_dimensions = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"+        strideAttr  = AttrRaw $ "window_strides = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> strides) <> ">"+        paddingAttr = AttrRaw $ "padding = dense<[["+            <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<"+            <> T.pack (show (length padding)) <> "x2xi64>"++    vid <- emitOpRegions "stablehlo.reduce_window"+            [input, initVal]+            [inType, initType]+            [windowAttr, strideAttr, paddingAttr]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)+  where+    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)+    padPair _      = error "breduceWindowAdd: padding must be [[low,high], ...]"++-- | Reduce a BTensor over specified window dimensions with @maximum@.+breduceWindowMax :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor+breduceWindowMax (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do+    let elemType = TensorType [] (ttDType inType)+    -- Build the max reduction region.+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @( 'F32)+        b <- arg @'[] @( 'F32)+        maxVid <- emitOp "stablehlo.maximum"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [maxVid] [elemType]++    let windowAttr  = AttrRaw $ "window_dimensions = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"+        strideAttr  = AttrRaw $ "window_strides = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> strides) <> ">"+        paddingAttr = AttrRaw $ "padding = dense<[["+            <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<"+            <> T.pack (show (length padding)) <> "x2xi64>"++    vid <- emitOpRegions "stablehlo.reduce_window"+            [input, initVal]+            [inType, initType]+            [windowAttr, strideAttr, paddingAttr]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)+  where+    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)+    padPair _      = error "breduceWindowMax: padding must be [[low,high], ...]"

@@ -13,6 +13,7 @@ import qualified Data.Text as T import qualified Data.Map.Strict as Map import Data.Map.Strict (Map)+import Debug.Trace (trace)  import HHLO.IR.AST import HHLO.IR.Builder@@ -61,6 +62,8 @@         "stablehlo.floor"        -> return cmap  -- non-differentiable         "stablehlo.ceil"         -> return cmap  -- non-differentiable         "stablehlo.sort"         -> error "autograd-hhlo: sort/topK is not differentiable"+        "stablehlo.reduce_window" -> vjpReduceWindow op resultBars cmap+        "stablehlo.convolution"   -> vjpConvolution op resultBars cmap         _ -> error $ T.unpack $ "autograd-hhlo: no VJP rule for " <> opName op   where     resultBars = map (\r -> Map.lookup r cmap) (opResults op)@@ -597,3 +600,354 @@         dx <- bmultiply bar oneMinus         accumulate cmap (btVid x) dx     Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- reduce_window VJP rule+-- ---------------------------------------------------------------------------++vjpReduceWindow :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindow op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            xType = btType x+            xVid  = btVid x+            xShape = ttShape xType+            rank = length xShape++        -- Detect reduction type by inspecting the region.+        let isAdd = any regionHasAdd (opRegions op)+            isMax = any regionHasMax (opRegions op)+            regionHasAdd (Region blocks) = any blockHasAdd blocks+            blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops+            regionHasMax (Region blocks) = any blockHasMax blocks+            blockHasMax (Block _ ops) = any (\o -> opName o == "stablehlo.maximum") ops++        -- Parse window attributes.+        let windowDims = findRawIntList "window_dimensions" (opAttributes op)+            strides    = findRawIntList "window_strides" (opAttributes op)+            padding    = findPadding (opAttributes op)++        if isAdd+            then vjpReduceWindowAdd bar xVid xType windowDims strides padding rank xShape cmap+            else if isMax+                then vjpReduceWindowMax bar x xType windowDims strides padding rank xShape cmap+                else error "autograd-hhlo: reduce_window VJP only supports add and maximum reductions"+    Nothing -> return cmap++vjpReduceWindowAdd :: BTensor -> ValueId -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindowAdd bar xVid xType windowDims strides padding _rank xShape cmap = do+    -- Only non-overlapping windows with VALID padding.+    let isNonOverlapping = and (zipWith (==) windowDims strides)+        isValid = all (all (== 0)) padding+    if not (isNonOverlapping && isValid)+        then error "autograd-hhlo: reduce_window(add) VJP only supports non-overlapping windows with VALID padding"+        else do+            -- For non-overlapping sum-pooling, each output gradient is+            -- broadcast back to its window.+            dx <- broadcastToInputShape bar xShape windowDims strides+            accumulate cmap xVid dx++vjpReduceWindowMax :: BTensor -> BTensor -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindowMax bar x xType windowDims strides padding _rank xShape cmap = do+    -- Only non-overlapping windows with VALID padding.+    let isNonOverlapping = and (zipWith (==) windowDims strides)+        isValid = all (all (== 0)) padding+    if not (isNonOverlapping && isValid)+        then error "autograd-hhlo: reduce_window(max) VJP only supports non-overlapping windows with VALID padding"+        else do+            let zeroValType = TensorType [] (ttDType xType)+                -- Compute the reduced output shape for NHWC non-overlapping pooling.+                outShape = map (\(sz, w) -> (sz - fromIntegral w) `div` fromIntegral w + 1) (zip xShape windowDims)+                outType = TensorType outShape (ttDType xType)+            -- Recompute forward max.+            negInf <- bconstant zeroValType (-1.0e30)+            maxVals <- breduceWindowMax x negInf windowDims strides padding outType+            -- Broadcast maxVals back to input shape.+            maxBroadcast <- broadcastToInputShape maxVals xShape windowDims strides+            -- Broadcast bar (gradient) back to input shape.+            barBroadcast <- broadcastToInputShape bar xShape windowDims strides+            -- mask = (input == maxBroadcast)+            mask <- bcompareEQ x maxBroadcast+            -- dx = select(mask, barBroadcast, 0)+            zero <- bconstant xType 0.0+            dx <- bselect mask barBroadcast zero xType+            accumulate cmap (btVid x) dx++-- | Broadcast a reduced tensor back to the original input shape for+-- non-overlapping reduce_window.+-- | Broadcast a reduced tensor back to the original input shape for+-- non-overlapping reduce_window (NHWC with 2 spatial dims).+broadcastToInputShape :: BTensor -> [Integer] -> [Int64] -> [Int64] -> Builder BTensor+broadcastToInputShape reduced xShape windowDims _strides = do+    let reducedShape = ttShape (btType reduced)+        -- reducedShape = [N, outH, outW, C]+        -- Insert size-1 after outH and outW.+        reshapedShape = [reducedShape !! 0, reducedShape !! 1, 1, reducedShape !! 2, 1, reducedShape !! 3]+        -- Broadcast to [N, outH, kh, outW, kw, C]+        broadcastShape = [reducedShape !! 0, reducedShape !! 1, fromIntegral (windowDims !! 1), reducedShape !! 2, fromIntegral (windowDims !! 2), reducedShape !! 3]+        broadcastDims = [0, 1, 2, 3, 4, 5] :: [Int64]+    reshaped <- breshape reduced (TensorType reshapedShape (ttDType (btType reduced)))+    broadcasted <- bbroadcastInDim reshaped broadcastDims (TensorType broadcastShape (ttDType (btType reduced)))+    -- Reshape back to [N, H, W, C]+    breshape broadcasted (TensorType xShape (ttDType (btType reduced)))++findRawIntList :: Text -> [Attribute] -> [Int64]+findRawIntList _ [] = []+findRawIntList name (AttrRaw raw : rest) =+    let key = name <> " = array<i64:"+    in if key `T.isInfixOf` raw+        then parseIntList raw+        else findRawIntList name rest+findRawIntList name (_ : rest) = findRawIntList name rest++findPadding :: [Attribute] -> [[Int64]]+findPadding [] = []+findPadding (AttrRaw raw : rest) =+    let key = "padding = dense<"+    in if key `T.isInfixOf` raw+        then parsePadding raw+        else findPadding rest+findPadding (_ : rest) = findPadding rest++parsePadding :: Text -> [[Int64]]+parsePadding raw =+    let inner = extractNestedBrackets raw+        pairs = T.splitOn "], [" inner+    in map parsePair pairs+  where+    extractNestedBrackets txt =+        let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt+        in go afterFirstBracket 1 ""+      where+        go txt' depth acc+            | T.null txt' = acc+            | T.head txt' == ']' && depth == 1 = acc+            | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[")+            | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]")+            | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt'])+    parsePair t =+        let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t))+        in map (read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Convolution VJP rule (NHWC only)+-- ---------------------------------------------------------------------------++vjpConvolution :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConvolution op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let input = operandBT op 0+            kernel = operandBT op 1+            inputType = btType input+            kernelType = btType kernel+            attrs = opAttributes op+            dimNums = lookupAttrString "dim_numbers" attrs+        -- Dispatch based on whether this is a regular conv or transpose conv.+        -- Regular conv:  "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"+        -- Transpose conv: "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"+        let windowStr = lookupAttrString "window" attrs+            (stride, pad, lhsDilate, _rhsDilate) = parseWindowString windowStr+        if dimNums == "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"+            then do+                -- Regular convolution.+                dInput <- convBackwardInput bar kernel kernelType inputType stride pad+                cmap' <- accumulate cmap (btVid input) dInput+                -- Only compute kernel gradient if the kernel is a function argument+                -- (negative ValueId) or if its gradient is already needed.+                let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap+                if needKernelGrad+                    then do+                        dKernel <- convBackwardKernel input inputType bar stride pad+                        accumulate cmap' (btVid kernel) dKernel+                    else return cmap'+            else if dimNums == "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"+            then do+                -- Transposed convolution.+                dInput <- transposeConvBackwardInput bar kernel inputType lhsDilate pad+                cmap' <- accumulate cmap (btVid input) dInput+                let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap+                if needKernelGrad+                    then do+                        dKernel <- transposeConvBackwardKernel input inputType bar lhsDilate pad+                        accumulate cmap' (btVid kernel) dKernel+                    else return cmap'+            else error "autograd-hhlo: convolution VJP only supports NHWC dim_numbers"+    Nothing -> return cmap++convBackwardInput :: BTensor -> BTensor -> TensorType -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+convBackwardInput bar kernel kernelType inputType stride pad = do+    -- Flip kernel spatially (dims 0 and 1).+    flippedKernel <- breverse kernel [0, 1] kernelType+    -- Backward input uses transposed conv dimension numbers.+    -- Window attributes only apply to spatial dims (first 2 of kernel shape).+    let spatialKernelShape = take 2 (ttShape kernelType)+        spatialReversePad = reversePad pad stride spatialKernelShape+        windowStr = buildWindowString [1, 1] spatialReversePad stride [1, 1]+    bconvolution bar flippedKernel "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType++convBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor+convBackwardKernel input inputType bar stride pad = do+    -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]+    let inputShape = ttShape inputType+        inputTShape = tail inputShape ++ [head inputShape]+    inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType))+    -- Transpose bar (dy): [N, outH, outW, C_out] -> [outH, outW, N, C_out]+    let barShape = ttShape (btType bar)+        barTShape = tail barShape ++ [head barShape]+    barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar)))+    -- Convolve with adapted dim numbers.  Window uses spatial dims only.+    let spatialKernelShape = take 2 (tail inputShape)+        outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType)+        windowStr = buildWindowString stride pad [1, 1] [1, 1]+    dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType+    -- Transpose output dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in]+    btranspose dk [0, 1, 3, 2] outType++-- ---------------------------------------------------------------------------+-- Transpose convolution backward helpers+-- ---------------------------------------------------------------------------++transposeConvBackwardInput :: BTensor -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+transposeConvBackwardInput bar kernel inputType lhsDilate pad = do+    -- Backward input: conv(dy, kernel) with stride = lhs_dilate.+    let windowStr = buildWindowString lhsDilate pad [1, 1] [1, 1]+    bconvolution bar kernel "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType++transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor+transposeConvBackwardKernel input inputType bar lhsDilate pad = do+    -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]+    let inputShape = ttShape inputType+        inputTShape = tail inputShape ++ [head inputShape]+    inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType))+    -- Transpose bar: [N, outH, outW, C_out] -> [outH, outW, N, C_out]+    let barShape = ttShape (btType bar)+        barTShape = tail barShape ++ [head barShape]+    barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar)))+    -- Convolve with lhs_dilate = forward_lhs_dilate.  Window uses spatial dims only.+    let spatialKernelShape = take 2 (tail inputShape)+        outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType)+        windowStr = buildWindowString [1, 1] pad lhsDilate [1, 1]+    dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType+    -- Transpose dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in]+    btranspose dk [0, 1, 3, 2] outType++-- ---------------------------------------------------------------------------+-- Shared attribute parsing helpers+-- ---------------------------------------------------------------------------++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++-- | Parse a plain bracketed int list like @[1, 2]@.+parsePlainIntList :: Text -> [Int64]+parsePlainIntList raw =+    let inner = T.takeWhile (/= ']') $ T.dropWhile (/= '[') raw+        withoutBracket = T.dropWhile (== '[') inner+        nums = T.splitOn "," withoutBracket+    in map (parseNum raw) nums+  where+    parseNum original t =+        let s = T.unpack (T.strip t)+        in case reads s of+            [(n, "")] -> fromIntegral (n :: Integer)+            _ -> error $ "parsePlainIntList: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack original ++ "'"++-- ---------------------------------------------------------------------------+-- Window attribute parsing helpers+-- ---------------------------------------------------------------------------++parseWindowString :: Text -> ([Int64], [[Int64]], [Int64], [Int64])+parseWindowString txt =+    let s = findField "stride" txt+        p = findField "pad" txt+        ld = findField "lhs_dilate" txt+        rd = findField "rhs_dilate" txt+    in ( parsePlainIntList s+       , parsePad p+       , if T.null ld then [1, 1] else parsePlainIntList ld+       , if T.null rd then [1, 1] else parsePlainIntList rd+       )+  where+    -- Find a field value, handling nested brackets for the 'pad' field.+    findField :: Text -> Text -> Text+    findField name t =+        let prefix = name <> " = "+        in case T.breakOn prefix t of+            (_, rest) | T.null rest -> ""+            (_, rest') ->+                let rest = T.drop (T.length prefix) rest'+                in takeValue rest++    -- Take the value, respecting bracket nesting.+    takeValue :: Text -> Text+    takeValue t = takeValueGo t (0 :: Int) ""+      where+        takeValueGo :: Text -> Int -> Text -> Text+        takeValueGo txt bracketDepth acc+            | T.null txt = acc+            | T.head txt == '}' && bracketDepth == 0 = acc+            | T.head txt == ',' && bracketDepth == 0 = acc+            | T.head txt == '[' = takeValueGo (T.tail txt) (bracketDepth + 1) (acc <> "[")+            | T.head txt == ']' = takeValueGo (T.tail txt) (bracketDepth - 1) (acc <> "]")+            | otherwise = takeValueGo (T.tail txt) bracketDepth (acc <> T.pack [T.head txt])++    parsePad :: Text -> [[Int64]]+    parsePad t+        | T.null t  = [[0, 0], [0, 0]]+        | otherwise =+            let -- Extract content between outermost [[ and ]]+                inner = extractNestedBrackets t+                pairs = T.splitOn "], [" inner+            in map parsePair pairs+      where+        extractNestedBrackets txt =+            let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt+            in go afterFirstBracket 1 ""+          where+            go txt' depth acc+                | T.null txt' = acc+                | T.head txt' == ']' && depth == 1 = acc+                | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[")+                | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]")+                | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt'])++    parsePair :: Text -> [Int64]+    parsePair t =+        let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t))+        in map parseNum nums+      where+        parseNum t' =+            let s = T.unpack (T.strip t')+            in case reads s of+                [(n, "")] -> fromIntegral (n :: Integer)+                _ -> error $ "parsePair: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack t ++ "'"++reversePad :: [[Int64]] -> [Int64] -> [Integer] -> [[Int64]]+reversePad pad _stride kernelShape =+    zipWith go pad (map fromIntegral kernelShape)+  where+    go [l, h] k = [fromIntegral (k :: Integer) - 1 - h, fromIntegral k - 1 - l]+    go _      _ = error "reversePad: invalid padding pair"++buildWindowString :: [Int64] -> [[Int64]] -> [Int64] -> [Int64] -> Text+buildWindowString stride pad lhsDilate rhsDilate =+    "{stride = " <> showList' stride+    <> ", pad = " <> showPad pad+    <> ", lhs_dilate = " <> showList' lhsDilate+    <> ", rhs_dilate = " <> showList' rhsDilate <> "}"+  where+    showList' xs = "[" <> T.intercalate ", " (map (T.pack . show) xs) <> "]"+    showPad ps = "[" <> T.intercalate ", " (map showPair ps) <> "]"+    showPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"+    showPair _      = error "buildWindowString: invalid padding pair"++lookupAttrString :: Text -> [Attribute] -> Text+lookupAttrString name = foldr f ""+  where+    f (AttrString n s) acc | n == name = s <> acc+    f _ acc = acc

@@ -164,6 +164,17 @@         fixPermAttr (AttrIntList "permutation" vals) =             AttrRaw $ "permutation = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"         fixPermAttr a = a+    pretty (Operation "stablehlo.reverse" operands operandTypes attrs regions results resultTypes) =+        -- Generic form with array<i64: ...> for dimensions (PJRT v1.16.0 compat).+        let attrs' = map fixRevAttr attrs+        in prettyResultVids results <> " = \"stablehlo.reverse\"("+           <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+           <> (if null attrs' then mempty else " " <> prettyAttrs attrs')+           <> " : " <> prettyResultType operandTypes resultTypes+      where+        fixRevAttr (AttrIntList "dimensions" vals) =+            AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"+        fixRevAttr a = a     pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions results resultTypes) =         -- Generic form (custom form syntax varies across parser versions).         prettyResultVids results <> " = \"stablehlo.concatenate\"("

@@ -46,4 +46,36 @@             modu = gradModule @'[2] @'F32 f             text = render modu         assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpReduceWindow (avgPool)" $ do+        let f x = do+                let windowDims = [1, 2, 2, 1]+                    strides    = [1, 2, 2, 1]+                    padding    = replicate 4 [0, 0]+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 1, 1, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            modu = gradModule @'[1, 4, 4, 1] @'F32 f+            text = render modu+        assertBool "contains reduce_window" ("reduce_window" `T.isInfixOf` text)+        assertBool "contains pad" ("pad" `T.isInfixOf` text)+    , testCase "vjpConvolution" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            modu = gradModule @'[1, 3, 3, 1] @'F32 f+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+        assertBool "contains reverse" ("reverse" `T.isInfixOf` text)+    , testCase "vjpTransposeConvolution" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @3 @3 [1, 2, 2, 1] (replicate 2 [0, 0]) x k+                sumAll y+            modu = gradModule @'[1, 2, 2, 1] @'F32 f+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)     ]

@@ -71,4 +71,62 @@         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)+    , testCase "grad avgPool" $ withPJRTCPU $ \api client -> do+        let f x = do+                let windowDims = [1, 2, 2, 1]+                    strides    = [1, 2, 2, 1]+                    padding    = replicate 4 [0, 0]+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 2, 2, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- grad = 1/4 for every element (non-overlapping 2x2 avg pool)+        let expected = V.fromList (replicate 16 0.25)+        assertBool "avgPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad conv2d" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            gradModu = gradModule @'[1, 3, 3, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..9.0]+        bufIn <- toDeviceF32 api client inp [1, 3, 3, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 9+        -- grad for 3x3 input with 2x2 kernel all 1s:+        -- corners: 1, edges: 2, center: 4+        let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]+        assertBool "conv2d grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad maxPool" $ withPJRTCPU $ \api client -> do+        let f x = do+                let kernel = [2, 2]+                    stride = [2, 2]+                    padding = [[0, 0], [0, 0]]+                y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- maxPool 2x2 stride 2 on 4x4:+        -- Window (0,0): max=6 at pos (1,1) -> index 5+        -- Window (0,1): max=8 at pos (1,3) -> index 7+        -- Window (1,0): max=14 at pos (3,1) -> index 13+        -- Window (1,1): max=16 at pos (3,3) -> index 15+        let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]+        assertBool "maxPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)     ]