hhlo 0.8.0.0 → 0.9.0.0

raw patch · 20 files changed

+292/−159 lines, 20 filesdep +vector-sizedPVP ok

version bump matches the API change (PVP)

Dependencies added: vector-sized

API changes (from Hackage documentation)

Files

• CHANGELOG.md +32/−1
• examples/22-new-ops-smoke-test.hs +2/−2
• examples/23-resnet.hs +15/−15
• examples/24-alexnet.hs +7/−7
• examples/26-unet.hs +15/−15
• hhlo.cabal +4/−2
• src/HHLO/Core/Types.hs +66/−0
• src/HHLO/EDSL/Ops.hs +104/−73
• test/Test/Autograd/Rules.hs +4/−4
• test/Test/EDSL/Ops.hs +11/−10
• test/Test/Runtime/EndToEndAutograd.hs +6/−6
• test/Test/Runtime/EndToEndAutogradGPU.hs +6/−6
• test/Test/Runtime/EndToEndDataMovement.hs +4/−4
• test/Test/Runtime/EndToEndDataMovementGPU.hs +4/−4
• test/Test/Runtime/EndToEndMatmul.hs +2/−1
• test/Test/Runtime/EndToEndMatmulGPU.hs +2/−1
• test/Test/Runtime/EndToEndReductions.hs +2/−2
• test/Test/Runtime/EndToEndReductionsGPU.hs +2/−2
• test/Test/Runtime/EndToEndShape.hs +2/−2
• test/Test/Runtime/EndToEndShapeGPU.hs +2/−2

@@ -162,4 +162,35 @@   (falling back to `deps/pjrt/`), so downstream libraries no longer need to   reimplement plugin discovery. -## next+## 0.9.0.0 -- 2026-04-20++* **Fixed-length configuration vectors** — rank-polymorphic EDSL ops now use+  `vector-sized` to tie config vector lengths to tensor ranks at compile time.+  This eliminates the class of bugs where wrong-length config silently produces+  invalid StableHLO.+  * Phase 1 (2D NN primitives): `conv2dWithPadding`, `maxPool`, `avgPool`,+    `transposeConvolution` accept `V2 Int64` / `P2`.+  * Phase 2+ (rank-polymorphic ops): `transpose`, `slice`, `pad`,+    `dynamicSlice`, `reduceWindow`, and `dotGeneral` now accept+    `Vector (Length s) Int64` or separate `Vector n Int64` type parameters+    instead of raw `[Int64]`.+  ```haskell+  -- BEFORE (could silently miscompile)+  transpose [1, 0] x+  slice x [1] [3] [1]+  dotGeneral [] [] [1] [0] x y++  -- AFTER (type-safe)+  transpose (v2 1 0) x+  slice x (v1 1) (v1 3) (v1 1)+  dotGeneral VS.empty VS.empty (v1 1) (v1 0) x y+  ```+  New exports in `HHLO.Core.Types`: `Length` type family, `V`, `V1`, `V2`,+  `V3`, `V4`, `Padding`, `P2`, plus smart constructors `v1`, `v2`, `v3`, `v4`,+  `p2`.+* New dependency: `vector-sized >= 1.5 && < 1.6`.+* Fix `transposeConvolution` lhs_dilation bug — passing a 2-element spatial+  dilation list no longer drops the second element.+* `gather` and `scatter` kept as `[Int64]` for now. Their config vector lengths+  depend on complex relationships between operand / indices / result ranks, so+  a clean type-safe design requires a separate future phase.

@@ -177,7 +177,7 @@             [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ]             $ do                 x <- arg @'[1, 4, 4, 1] @'F32-                maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x     exec <- compile api client (render modu)     let input = V.fromList             [ 1,  2,  3,  4@@ -197,7 +197,7 @@             [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ]             $ do                 x <- arg @'[1, 4, 4, 1] @'F32-                avgPool [2, 2] [2, 2] x+                avgPool (v2 2 2) (v2 2 2) x     exec <- compile api client (render modu)     let input = V.fromList             [ 1,  2,  3,  4

@@ -40,25 +40,25 @@ initialConv :: Tensor '[1, 8, 8, 3] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) initialConv x = do     w <- constant @'[3, 3, 3, 16] @'F32 0.01-    conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w+    conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w  -- Stage 1: 2 basic blocks, 16 channels, 8x8 (no downsampling) stage1Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) stage1Block1 x = do     w1 <- constant @'[3, 3, 16, 16] @'F32 0.01-    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 16, 16] @'F32 0.01-    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  stage1Block2 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) stage1Block2 x = do     w1 <- constant @'[3, 3, 16, 16] @'F32 0.01-    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 16, 16] @'F32 0.01-    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  -- Stage 2: 2 blocks, 16 -> 32 channels, 8x8 -> 4x4 (stride 2 on first)@@ -66,43 +66,43 @@ stage2Block1 x = do     -- Projection shortcut     wSkip <- constant @'[1, 1, 16, 32] @'F32 0.01-    skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 [2, 2] [[0, 0], [0, 0]] x wSkip+    skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 (v2 2 2) (p2 (0,0) (0,0)) x wSkip     -- Main path     w1 <- constant @'[3, 3, 16, 32] @'F32 0.01-    out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 [2, 2] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 (v2 2 2) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 32, 32] @'F32 0.01-    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out skip  stage2Block2 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32) stage2Block2 x = do     w1 <- constant @'[3, 3, 32, 32] @'F32 0.01-    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 32, 32] @'F32 0.01-    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  -- Stage 3: 2 blocks, 32 -> 64 channels, 4x4 -> 2x2 (stride 2 on first) stage3Block1 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) stage3Block1 x = do     wSkip <- constant @'[1, 1, 32, 64] @'F32 0.01-    skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 [2, 2] [[0, 0], [0, 0]] x wSkip+    skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x wSkip     w1 <- constant @'[3, 3, 32, 64] @'F32 0.01-    out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 [2, 2] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 64, 64] @'F32 0.01-    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out skip  stage3Block2 :: Tensor '[1, 2, 2, 64] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) stage3Block2 x = do     w1 <- constant @'[3, 3, 64, 64] @'F32 0.01-    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] x w1+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) x w1     out <- relu out     w2 <- constant @'[3, 3, 64, 64] @'F32 0.01-    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  -- ---------------------------------------------------------------------------

@@ -53,29 +53,29 @@                  -- Conv 3x3/1, 32 channels                 w1 <- constant @'[3, 3, 3, 32] @'F32 0.01-                x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x w1+                x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) x w1                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x                  -- Conv 3x3/1, 64 channels                 w2 <- constant @'[3, 3, 32, 64] @'F32 0.01-                x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w2+                x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w2                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x                  -- Conv 3x3/1, 128 channels                 w3 <- constant @'[3, 3, 64, 128] @'F32 0.01-                x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w3+                x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w3                 x <- relu x                  -- Conv 3x3/1, 128 channels                 w4 <- constant @'[3, 3, 128, 128] @'F32 0.01-                x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w4+                x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w4                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x                 let x' :: Tensor '[1, 2, 2, 128] 'F32                     x' = x 

@@ -55,34 +55,34 @@                  -- Encoder stage 1: 1 -> 16 channels                 wE1a <- constant @'[3, 3, 1, 16] @'F32 0.01-                e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x wE1a+                e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) x wE1a                 e1 <- relu e1                 wE1b <- constant @'[3, 3, 16, 16] @'F32 0.01-                e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] e1 wE1b+                e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) e1 wE1b                 e1 <- relu e1                 -- Skip1 = e1 (shape [1,16,16,16])                  -- Downsample: maxPool 2x2/2-                d1 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e1+                d1 <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) e1                  -- Encoder stage 2: 16 -> 32 channels                 wE2a <- constant @'[3, 3, 16, 32] @'F32 0.01-                e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] d1 wE2a+                e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) d1 wE2a                 e2 <- relu e2                 wE2b <- constant @'[3, 3, 32, 32] @'F32 0.01-                e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] e2 wE2b+                e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) e2 wE2b                 e2 <- relu e2                 -- Skip2 = e2 (shape [1,8,8,32])                  -- Downsample: maxPool 2x2/2-                d2 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e2+                d2 <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) e2                  -- ========== Bottleneck ==========                 wBa <- constant @'[3, 3, 32, 64] @'F32 0.01-                b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] d2 wBa+                b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) d2 wBa                 b <- relu b                 wBb <- constant @'[3, 3, 64, 64] @'F32 0.01-                b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] b wBb+                b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) b wBb                 b <- relu b                  -- ========== Decoder ==========@@ -90,31 +90,31 @@                 -- Decoder stage 2: upsample 4x4 -> 8x8, concat with skip2                 -- Transpose conv: [1,4,4,64] -> [1,8,8,32]                 wD2up <- constant @'[2, 2, 32, 64] @'F32 0.01-                u2 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] b wD2up+                u2 <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) b wD2up                 -- Concatenate with skip2 along channel dim (axis 3)                 u2 <- concatenate2 @'[1, 8, 8, 32] @'[1, 8, 8, 32] @'[1, 8, 8, 64] @'F32 3 u2 e2                 wD2a <- constant @'[3, 3, 64, 32] @'F32 0.01-                u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2a+                u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) u2 wD2a                 u2 <- relu u2                 wD2b <- constant @'[3, 3, 32, 32] @'F32 0.01-                u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2b+                u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) u2 wD2b                 u2 <- relu u2                  -- Decoder stage 1: upsample 8x8 -> 16x16, concat with skip1                 wD1up <- constant @'[2, 2, 16, 32] @'F32 0.01-                u1 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] u2 wD1up+                u1 <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) u2 wD1up                 -- Concatenate with skip1 along channel dim (axis 3)                 u1 <- concatenate2 @'[1, 16, 16, 16] @'[1, 16, 16, 16] @'[1, 16, 16, 32] @'F32 3 u1 e1                 wD1a <- constant @'[3, 3, 32, 16] @'F32 0.01-                u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1a+                u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) u1 wD1a                 u1 <- relu u1                 wD1b <- constant @'[3, 3, 16, 16] @'F32 0.01-                u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1b+                u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) u1 wD1b                 u1 <- relu u1                  -- Final 1x1 conv: 16 -> 2 channels                 wOut <- constant @'[1, 1, 16, 2] @'F32 0.01-                conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 [1, 1] [[0, 0], [0, 0]] u1 wOut+                conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 (v2 1 1) (p2 (0,0) (0,0)) u1 wOut      putStrLn "Generated MLIR (first 20 lines):"     let lines_ = T.lines (render modu)

@@ -1,7 +1,7 @@ cabal-version:      3.0 name:               hhlo-version:            0.8.0.0-synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU+version:            0.9.0.0+synopsis:           Haskell Frontend for StableHLO — type-safe ML training/inference on CPU and GPU description:     HHLO is a Haskell library and runtime for building, compiling, and executing     machine-learning programs targeting StableHLO, the portable intermediate@@ -82,6 +82,7 @@         text          >= 2.0   && < 2.2,         bytestring    >= 0.12  && < 0.13,         vector        >= 0.13  && < 0.14,+        vector-sized  >= 1.5   && < 1.6,         containers    >= 0.6   && < 0.8,         transformers  >= 0.6   && < 0.7,         mtl           >= 2.3   && < 2.4,@@ -528,6 +529,7 @@         text          >= 2.0   && < 2.2,         bytestring    >= 0.12  && < 0.13,         vector        >= 0.13  && < 0.14,+        vector-sized  >= 1.5   && < 1.6,         tasty         >= 1.5   && < 1.6,         tasty-hunit   >= 0.10  && < 0.11     hs-source-dirs:   test

@@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE KindSignatures #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ScopedTypeVariables #-}  {-# LANGUAGE OverloadedStrings #-}@@ -12,6 +13,20 @@     , KnownShape(..)     , dtypeToText     , HostType+    -- * Fixed-length configuration vectors+    , Length+    , V+    , V1+    , V2+    , V3+    , V4+    , Padding+    , P2+    , v1+    , v2+    , v3+    , v4+    , p2     ) where  import GHC.TypeLits@@ -20,6 +35,7 @@ import Data.Int (Int8, Int16, Int32, Int64) import Data.Word (Word8, Word16, Word32, Word64) import Data.Text (Text)+import qualified Data.Vector.Sized as VS   -- | Supported element types for tensors.@@ -79,3 +95,53 @@     HostType 'UI32 = Word32     HostType 'UI64 = Word64     HostType 'Bool = Word8++-- ---------------------------------------------------------------------------+-- Fixed-length configuration vectors+-- ---------------------------------------------------------------------------++-- | Compute the length of a type-level list.+type family Length (xs :: [k]) :: Nat where+    Length '[]     = 0+    Length (x:xs)  = 1 + Length xs++-- | Fixed-length vector alias.+type V (n :: Nat) a = VS.Vector n a++-- | 1-element vector.+type V1 a = V 1 a++-- | 2-element vector (e.g. spatial height, width).+type V2 a = V 2 a++-- | 3-element vector.+type V3 a = V 3 a++-- | 4-element vector (e.g. NHWC dimensions).+type V4 a = V 4 a++-- | Padding config: one (low,high) pair per dimension.+type Padding (n :: Nat) = V n (Int64, Int64)++-- | 2D padding alias (common case).+type P2 = Padding 2++-- | Smart constructor for a 1-element vector.+v1 :: a -> V1 a+v1 a = a `VS.cons` VS.empty++-- | Smart constructor for a 2-element vector.+v2 :: a -> a -> V2 a+v2 a b = a `VS.cons` (b `VS.cons` VS.empty)++-- | Smart constructor for a 3-element vector.+v3 :: a -> a -> a -> V3 a+v3 a b c = a `VS.cons` (b `VS.cons` (c `VS.cons` VS.empty))++-- | Smart constructor for a 4-element vector.+v4 :: a -> a -> a -> a -> V4 a+v4 a b c d = a `VS.cons` (b `VS.cons` (c `VS.cons` (d `VS.cons` VS.empty)))++-- | Smart constructor for 2D padding from two (before,after) pairs.+p2 :: (Int64, Int64) -> (Int64, Int64) -> P2+p2 = v2

@@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-}+{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE OverloadedStrings #-} @@ -60,6 +61,19 @@     , conv2dWithPadding     , transposeConvolution     , batchNormInference+    -- * Fixed-length configuration helpers+    , Length+    , V+    , V1+    , V2+    , V3+    , V4+    , Padding+    , v1+    , v2+    , v3+    , v4+    , p2     , layerNorm     , globalAvgPool     , gelu@@ -148,7 +162,21 @@ import HHLO.Core.Types import HHLO.IR.AST import HHLO.IR.Builder+import qualified Data.Vector.Sized as VS+import Data.Vector.Sized (Vector)+import qualified Data.Vector as V+import Unsafe.Coerce (unsafeCoerce) +-- | Unsafe helper: convert a list to a fixed-length vector.+-- The caller must ensure the list length equals the rank of shape @s@.+vecFromListUnsafe :: forall s a. KnownShape s => [a] -> Vector (Length s) a+vecFromListUnsafe xs =+    let expected = length (shapeVal (Proxy @s))+        actual   = length xs+    in if expected == actual+       then unsafeCoerce (V.fromList xs)+       else error $ "vecFromListUnsafe: expected length " ++ show expected ++ ", got " ++ show actual+ -- --------------------------------------------------------------------------- -- Binary element-wise ops -- ---------------------------------------------------------------------------@@ -225,12 +253,12 @@  -- | General dot product with explicit batch and contracting dimensions. -- Use this for batched matrix multiplication (rank > 2).-dotGeneral :: forall s1 s2 sOut d.-              (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d)-           => [Int64]   -- ^ lhs batch dims-           -> [Int64]   -- ^ rhs batch dims-           -> [Int64]   -- ^ lhs contracting dims-           -> [Int64]   -- ^ rhs contracting dims+dotGeneral :: forall s1 s2 sOut d nBatch nContract.+              (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d, KnownNat nBatch, KnownNat nContract)+           => Vector nBatch Int64     -- ^ lhs batch dims+           -> Vector nBatch Int64     -- ^ rhs batch dims+           -> Vector nContract Int64  -- ^ lhs contracting dims+           -> Vector nContract Int64  -- ^ rhs contracting dims            -> Tensor s1 d            -> Tensor s2 d            -> Builder (Tensor sOut d)@@ -238,8 +266,8 @@     let 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) <> "]")+        batchAttr      = AttrString "batching_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList lhsBatch)) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList rhsBatch)) <> "]")+        contractingAttr = AttrString "contracting_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList lhsContract)) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList rhsContract)) <> "]")     vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]             [ batchAttr             , contractingAttr@@ -343,11 +371,11 @@ -- -- @perm@ must be a permutation of @[0 .. rank-1]@. transpose :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d)-          => [Int64] -> Tensor sIn d -> Builder (Tensor sOut d)+          => Vector (Length sIn) Int64 -> Tensor sIn d -> Builder (Tensor sOut d) transpose perm (Tensor x) = do     let inType  = tensorType (Proxy @sIn)  (Proxy @d)         outType = tensorType (Proxy @sOut) (Proxy @d)-        permAttr = AttrIntList "permutation" perm+        permAttr = AttrIntList "permutation" (VS.toList perm)     vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType     return (Tensor vid) @@ -518,7 +546,7 @@        => Tensor '[batch, h, w, inCh] 'F32        -> Tensor '[kh, kw, inCh, outCh] 'F32        -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)-conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow [1, 1] (replicate 2 [0, 0]) x k+conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow (v2 1 1) (p2 (0, 0) (0, 0)) x k  -- | 2-D convolution with explicit stride and padding (NHWC format). --@@ -530,8 +558,8 @@ conv2dWithPadding :: forall batch h w inCh outCh kh kw oh ow.                      ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh                      , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow )-                  => [Int64]        -- ^ strides [sh, sw]-                  -> [[Int64]]      -- ^ padding [[pt, pb], [pl, pr]]+                  => V2 Int64    -- ^ strides [sh, sw]+                  -> P2          -- ^ padding [[pt, pb], [pl, pr]]                   -> Tensor '[batch, h, w, inCh] 'F32                   -> Tensor '[kh, kw, inCh, outCh] 'F32                   -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)@@ -540,8 +568,8 @@         inType2 = tensorType (Proxy @'[kh, kw, inCh, outCh]) (Proxy @'F32)         outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32)         dimNums = "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"-        strideStr = "[" <> T.intercalate ", " ((T.pack . show) <$> strides) <> "]"-        padStr = "[" <> T.intercalate ", " (padPair <$> padding) <> "]"+        strideStr = "[" <> T.intercalate ", " ((T.pack . show) <$> VS.toList strides) <> "]"+        padStr = "[" <> padPair (padding `VS.index` 0) <> ", " <> padPair (padding `VS.index` 1) <> "]"         window  = "{stride = " <> strideStr <> ", pad = " <> padStr <> "}"     vid <- emitOp "stablehlo.convolution"             [tensorValue input, tensorValue kernel]@@ -553,8 +581,7 @@             ] outType     return (Tensor vid)   where-    padPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"-    padPair _      = error "conv2dWithPadding: padding must be [[low,high], ...]"+    padPair (l, h) = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"  -- | Batch normalization for inference. --@@ -768,11 +795,11 @@           , KnownShape sResult, KnownDType d )        => Tensor sOperand d        -> Tensor sIndices 'I64-       -> [Int64]   -- ^ offset_dims-       -> [Int64]   -- ^ collapsed_slice_dims-       -> [Int64]   -- ^ start_index_map-       -> Int64     -- ^ index_vector_dim-       -> [Int64]   -- ^ slice_sizes+       -> [Int64]      -- ^ offset_dims+       -> [Int64]     -- ^ collapsed_slice_dims+       -> [Int64]     -- ^ start_index_map+       -> Int64                               -- ^ index_vector_dim+       -> [Int64]     -- ^ slice_sizes        -> Builder (Tensor sResult d) gather operand indices offsetDims collapsedSliceDims startIndexMap indexVectorDim sliceSizes = do     let operandType = tensorType (Proxy @sOperand) (Proxy @d)@@ -813,10 +840,10 @@         -> Tensor sIndices 'I64         -> Tensor sUpdates d         -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] d))-        -> [Int64]   -- ^ update_window_dims-        -> [Int64]   -- ^ inserted_window_dims-        -> [Int64]   -- ^ scatter_dims_to_operand_dims-        -> Int64     -- ^ index_vector_dim+        -> [Int64]     -- ^ update_window_dims+        -> [Int64]       -- ^ inserted_window_dims+        -> [Int64]     -- ^ scatter_dims_to_operand_dims+        -> Int64                               -- ^ index_vector_dim         -> Builder (Tensor sResult d) scatter input indices updates updateFn updateWindowDims insertedWindowDims scatterDimsToOperandDims indexVectorDim = do     let inputType   = tensorType (Proxy @sInput)   (Proxy @d)@@ -860,16 +887,16 @@ slice :: forall sIn sOut d.          (KnownShape sIn, KnownShape sOut, KnownDType d)       => Tensor sIn d-      -> [Int64]   -- ^ start_indices-      -> [Int64]   -- ^ limit_indices-      -> [Int64]   -- ^ strides+      -> Vector (Length sIn) Int64   -- ^ start_indices+      -> Vector (Length sIn) Int64   -- ^ limit_indices+      -> Vector (Length sIn) Int64   -- ^ strides       -> Builder (Tensor sOut d) slice operand start limit stride = do     let inType   = tensorType (Proxy @sIn)  (Proxy @d)         outType  = tensorType (Proxy @sOut) (Proxy @d)-        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) <> ">"+        startAttr = AttrRaw $ "start_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList start) <> ">"+        limitAttr = AttrRaw $ "limit_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList limit) <> ">"+        strideAttr = AttrRaw $ "strides = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList stride) <> ">"      let (Tensor operandVid) = operand     vid <- emitOp "stablehlo.slice" [operandVid] [inType]@@ -884,17 +911,17 @@        (KnownShape sIn, KnownShape sOut, KnownDType d)     => Tensor sIn d     -> Tensor '[] d   -- ^ padding value (scalar)-    -> [Int64]        -- ^ edge_padding_low-    -> [Int64]        -- ^ edge_padding_high-    -> [Int64]        -- ^ interior_padding+    -> Vector (Length sIn) Int64   -- ^ edge_padding_low+    -> Vector (Length sIn) Int64   -- ^ edge_padding_high+    -> Vector (Length sIn) Int64   -- ^ interior_padding     -> Builder (Tensor sOut d) pad operand paddingValue low high interior = do     let inType   = tensorType (Proxy @sIn)  (Proxy @d)         padType  = tensorType (Proxy @'[])  (Proxy @d)         outType  = tensorType (Proxy @sOut) (Proxy @d)-        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) <> ">"+        lowAttr  = AttrRaw $ "edge_padding_low = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList low) <> ">"+        highAttr = AttrRaw $ "edge_padding_high = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList high) <> ">"+        intAttr  = AttrRaw $ "interior_padding = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList interior) <> ">"      let (Tensor operandVid) = operand         (Tensor padVid)     = paddingValue@@ -907,12 +934,12 @@                 (KnownShape sIn, KnownShape sOut, KnownDType d)              => Tensor sIn d              -> [Tensor '[] 'I64]   -- ^ start indices (one scalar i64 per dimension)-             -> [Int64]             -- ^ slice_sizes+             -> Vector (Length sOut) Int64   -- ^ slice_sizes              -> Builder (Tensor sOut d) dynamicSlice operand startIndices sliceSizes = do     let inType   = tensorType (Proxy @sIn)  (Proxy @d)         outType  = tensorType (Proxy @sOut) (Proxy @d)-        sizesAttr = AttrRaw $ "slice_sizes = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">"+        sizesAttr = AttrRaw $ "slice_sizes = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList sliceSizes) <> ">"      let (Tensor operandVid) = operand         startVids = tensorValue <$> startIndices@@ -1222,9 +1249,9 @@ -- @reduction = "stablehlo.add"@ and an init value of @0@. reduceWindow :: forall sIn sOut d.                 (KnownShape sIn, KnownShape sOut, KnownDType d)-             => [Int64]        -- ^ window_dimensions-             -> [Int64]        -- ^ window_strides-             -> [[Int64]]      -- ^ padding: [[low, high], ...] per dimension+             => Vector (Length sIn) Int64        -- ^ window_dimensions+             -> Vector (Length sIn) Int64        -- ^ window_strides+             -> Vector (Length sIn) (Int64, Int64) -- ^ padding: (low, high) per dimension              -> Text           -- ^ reduction op, e.g. "stablehlo.maximum"              -> Tensor '[] d   -- ^ init value (scalar)              -> Tensor sIn d   -- ^ input@@ -1247,12 +1274,12 @@         emitReturn [tensorValue result] [elemType]      let windowAttr = AttrRaw $ "window_dimensions = array<i64: "-            <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"+            <> T.intercalate ", " ((T.pack . show) <$> VS.toList 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>"+            <> T.intercalate ", " ((T.pack . show) <$> VS.toList strides) <> ">"+        padStr = "[" <> T.intercalate ", " (padPairV <$> VS.toList padding) <> "]"+        paddingAttr = AttrRaw $ "padding = dense<" <> padStr <> "> : tensor<"+            <> T.pack (show (length (shapeVal (Proxy @sIn)))) <> "x2xi64>"      let (Tensor initVid) = initVal         (Tensor inputVid) = input@@ -1265,21 +1292,23 @@             outType     return (Tensor vid)   where-    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)-    padPair _      = error "reduceWindow: padding must be [[low,high], ...]"+    padPairV (l, h) = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"  -- | 2-D max pooling (NHWC). maxPool :: forall n h w c oh ow.            (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow)-        => [Int64]   -- ^ kernel [kh, kw]-        -> [Int64]   -- ^ stride [sh, sw]-        -> [[Int64]] -- ^ padding per spatial dim [[pt, pb], [pl, pr]]+        => V2 Int64   -- ^ kernel [kh, kw]+        -> V2 Int64   -- ^ stride [sh, sw]+        -> P2         -- ^ padding per spatial dim [[pt, pb], [pl, pr]]         -> Tensor '[n, h, w, c] 'F32         -> Builder (Tensor '[n, oh, ow, c] 'F32) maxPool kernel stride padding x = do-    let windowDims = [1, kernel !! 0, kernel !! 1, 1]-        strides    = [1, stride !! 0, stride !! 1, 1]-        fullPadding = [[0, 0], padding !! 0, padding !! 1, [0, 0]]+    let windowDims = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1+        strides    = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1+        fullPadding = v4 (0, 0)+                         (fst (padding `VS.index` 0), snd (padding `VS.index` 0))+                         (fst (padding `VS.index` 1), snd (padding `VS.index` 1))+                         (0, 0)     -- init value for max: a very negative number     initVal <- constant @'[] @'F32 (-1.0e30)     reduceWindow windowDims strides fullPadding "stablehlo.maximum" initVal x@@ -1290,15 +1319,15 @@ -- by the input size, kernel, and stride (no padding). avgPool :: forall n h w c oh ow.            (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow)-        => [Int64]   -- ^ kernel [kh, kw]-        -> [Int64]   -- ^ stride [sh, sw]+        => V2 Int64   -- ^ kernel [kh, kw]+        -> V2 Int64   -- ^ stride [sh, sw]         -> Tensor '[n, h, w, c] 'F32         -> Builder (Tensor '[n, oh, ow, c] 'F32) avgPool kernel stride x = do-    let windowDims = [1, kernel !! 0, kernel !! 1, 1]-        strides    = [1, stride !! 0, stride !! 1, 1]-        fullPadding = replicate 4 [0, 0]-        windowSize = fromIntegral (product kernel) :: Double+    let windowDims = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1+        strides    = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1+        fullPadding = v4 (0, 0) (0, 0) (0, 0) (0, 0)+        windowSize = fromIntegral (product (VS.toList kernel)) :: Double     initVal <- constant @'[] @'F32 0.0     summed <- reduceWindow windowDims strides fullPadding "stablehlo.add" initVal x     divisor <- constant @'[] @'F32 windowSize@@ -1318,8 +1347,8 @@                         ( KnownNat batch, KnownNat h, KnownNat w                         , KnownNat inCh, KnownNat outCh                         , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow )-                     => [Int64]   -- ^ lhs_dilation (upsample factor), e.g. [1,2,2,1]-                     -> [[Int64]] -- ^ padding per spatial dim, e.g. [[1,1],[1,1]] for 2x2 kernel+                     => V2 Int64   -- ^ lhs_dilation (upsample factor), e.g. [2,2]+                     -> P2         -- ^ padding per spatial dim, e.g. p2 (1,1) (1,1) for 2x2 kernel                      -> Tensor '[batch, h, w, inCh] 'F32                      -> Tensor '[kh, kw, outCh, inCh] 'F32                      -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)@@ -1328,10 +1357,10 @@         inType2 = tensorType (Proxy @'[kh, kw, outCh, inCh]) (Proxy @'F32)         outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32)         dimNums = "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"-        padStr  = "[[" <> T.pack (show (padding !! 0 !! 0)) <> ", " <> T.pack (show (padding !! 0 !! 1)) <> "], ["-               <> T.pack (show (padding !! 1 !! 0)) <> ", " <> T.pack (show (padding !! 1 !! 1)) <> "]]"+        padStr  = "[[" <> T.pack (show (fst (padding `VS.index` 0))) <> ", " <> T.pack (show (snd (padding `VS.index` 0))) <> "], ["+               <> T.pack (show (fst (padding `VS.index` 1))) <> ", " <> T.pack (show (snd (padding `VS.index` 1))) <> "]]"         window  = "{stride = [1, 1], pad = " <> padStr-               <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> drop 1 (take 3 lhsDilation)) <> "]"+               <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> VS.toList lhsDilation) <> "]"                <> ", rhs_dilate = [1, 1]}"     vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel]             [inType1, inType2]@@ -1869,7 +1898,7 @@ slice1 :: forall n d. (KnownShape '[n], KnownDType d)        => Tensor '[n] d -> Int64 -> Builder (Tensor '[] d) slice1 vec i = do-    sliced <- slice @'[n] @'[1] @d vec [i] [i + 1] [1]+    sliced <- slice @'[n] @'[1] @d vec (v1 i) (v1 (i + 1)) (v1 1)     reshape @'[1] @'[] sliced  -- | Pack two scalar tensors into a rank-1 tensor of shape @[2]@.@@ -1933,7 +1962,7 @@ -- | 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)+         (KnownShape sIn, KnownShape sOut, KnownDType d, KnownNat (Length sIn))       => Int64            -- ^ dimension to split along       -> Int64            -- ^ number of equal splits       -> Tensor sIn d@@ -1943,12 +1972,14 @@         rank = length sInShape         dimSize = fromIntegral (sInShape !! fromIntegral dim) :: Int         chunkSize = dimSize `div` fromIntegral n-        stride = replicate rank 1+        stride = VS.replicate 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]]+        let startList = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]]+            limitList = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]]+            start = fromJust (VS.fromList startList)+            limit = fromJust (VS.fromList limitList)         slice @sIn @sOut @d t start limit stride         ) [0 .. fromIntegral n - 1] @@ -2008,7 +2039,7 @@     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+    slice @s @sOut @d sorted (vecFromListUnsafe @s start) (vecFromListUnsafe @s limit) (vecFromListUnsafe @s stride)  -- --------------------------------------------------------------------------- -- Einsum

@@ -47,9 +47,9 @@         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]+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0,0) (0,0) (0,0) (0,0)                 initVal <- constant @'[] @'F32 0.0                 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x                 divisor <- constant @'[] @'F32 4.0@@ -72,7 +72,7 @@     , 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+                y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 2 2) (p2 (0,0) (0,0)) x k                 sumAll y             modu = gradModule @'[1, 2, 2, 1] @'F32 f             text = render modu

@@ -6,6 +6,7 @@  import Prelude hiding (map, maximum, minimum, negate, compare, tanh, sqrt, sin, cos, tan, floor) import qualified Data.Text as T+import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -80,7 +81,7 @@                     [ FuncArg "arg0" (TensorType [3, 2] F32) ]                     $ do                         x <- arg @'[3, 2] @'F32-                        y <- transpose @'[3, 2] @'[2, 3] [1, 0] x+                        y <- transpose @'[3, 2] @'[2, 3] (v2 1 0) x                         return y             let rendered = render modu             assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered@@ -147,7 +148,7 @@                     $ do                         x <- arg @'[2, 3] @'F32                         y <- arg @'[3, 2] @'F32-                        z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 [] [] [1] [0] x y+                        z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 VS.empty VS.empty (v1 1) (v1 0) x y                         return z             let rendered = render modu             assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered@@ -181,7 +182,7 @@                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32                         initVal <- constant @'[] @'F32 0.0-                        y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] [1, 2, 2, 1] [1, 2, 2, 1] [[0, 0], [0, 0], [0, 0], [0, 0]] "stablehlo.add" initVal x+                        y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] (v4 1 2 2 1) (v4 1 2 2 1) (v4 (0,0) (0,0) (0,0) (0,0)) "stablehlo.add" initVal x                         return y             let rendered = render modu             assertBool "stablehlo.reduce_window" $ "stablehlo.reduce_window" `T.isInfixOf` rendered@@ -190,7 +191,7 @@                     [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32-                        y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                        y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x                         return y             let rendered = render modu             assertBool "reduce_window for maxPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered@@ -200,7 +201,7 @@                     [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32-                        y <- avgPool [2, 2] [2, 2] x+                        y <- avgPool (v2 2 2) (v2 2 2) x                         return y             let rendered = render modu             assertBool "reduce_window for avgPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered@@ -232,7 +233,7 @@                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32                         k <- constant @'[3, 3, 1, 1] @'F32 0.5-                        y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x k+                        y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x k                         return y             let rendered = render modu             assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered@@ -295,7 +296,7 @@                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32                         k <- constant @'[2, 2, 1, 1] @'F32 0.5-                        y <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] x k+                        y <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) x k                         return y             let rendered = render modu             assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered@@ -356,7 +357,7 @@                     [ FuncArg "arg0" (TensorType [4] F32) ]                     $ do                         x <- arg @'[4] @'F32-                        y <- slice x [1] [3] [1]+                        y <- slice x (v1 1) (v1 3) (v1 1)                         return y             let rendered = render modu             assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered@@ -366,7 +367,7 @@                     $ do                         x <- arg @'[2] @'F32                         padVal <- constant @'[] @'F32 0.0-                        y <- pad x padVal [1] [1] [0]+                        y <- pad x padVal (v1 1) (v1 1) (v1 0)                         return y             let rendered = render modu             assertBool "stablehlo.pad" $ "stablehlo.pad" `T.isInfixOf` rendered@@ -376,7 +377,7 @@                     $ do                         x <- arg @'[4] @'F32                         idx <- constant @'[] @'I64 1-                        y <- dynamicSlice x [idx] [2]+                        y <- dynamicSlice x [idx] (v1 2)                         return y             let rendered = render modu             assertBool "stablehlo.dynamic_slice" $ "stablehlo.dynamic_slice" `T.isInfixOf` rendered

@@ -79,9 +79,9 @@             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]+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0, 0) (0, 0) (0, 0) (0, 0)                 initVal <- constant @'[] @'F32 0.0                 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x                 divisor <- constant @'[] @'F32 4.0@@ -116,9 +116,9 @@             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]]+                let kernel = v2 2 2+                    stride = v2 2 2+                    padding = p2 (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

@@ -80,9 +80,9 @@     , testCase "grad avgPool" $ do         GPUResource api client dev <- getGPU         let f x = do-                let windowDims = [1, 2, 2, 1]-                    strides    = [1, 2, 2, 1]-                    padding    = replicate 4 [0, 0]+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0, 0) (0, 0) (0, 0) (0, 0)                 initVal <- constant @'[] @'F32 0.0                 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x                 divisor <- constant @'[] @'F32 4.0@@ -116,9 +116,9 @@     , testCase "grad maxPool" $ do         GPUResource api client dev <- getGPU         let f x = do-                let kernel = [2, 2]-                    stride = [2, 2]-                    padding = [[0, 0], [0, 0]]+                let kernel = v2 2 2+                    stride = v2 2 2+                    padding = p2 (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

@@ -28,7 +28,7 @@                 [ FuncArg "arg0" (TensorType [5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [1] [4] [1]+                    y <- slice x (v1 1) (v1 4) (v1 1)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -41,7 +41,7 @@                 [ FuncArg "arg0" (TensorType [5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [0] [4] [2]+                    y <- slice x (v1 0) (v1 4) (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -55,7 +55,7 @@                 $ do                     x <- arg @'[2] @'F32                     padVal <- constant @'[] @'F32 0.0-                    y <- pad x padVal [1] [1] [0]+                    y <- pad x padVal (v1 1) (v1 1) (v1 0)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [1.0, 2.0]@@ -200,7 +200,7 @@                 $ do                     x <- arg @'[4] @'F32                     idx <- constant @'[] @'I64 1-                    y <- dynamicSlice x [idx] [2]+                    y <- dynamicSlice x [idx] (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0]

@@ -30,7 +30,7 @@                 [ FuncArg "arg0" (TensorType [5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [1] [4] [1]+                    y <- slice x (v1 1) (v1 4) (v1 1)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -44,7 +44,7 @@                 [ FuncArg "arg0" (TensorType [5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [0] [4] [2]+                    y <- slice x (v1 0) (v1 4) (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -59,7 +59,7 @@                 $ do                     x <- arg @'[2] @'F32                     padVal <- constant @'[] @'F32 0.0-                    y <- pad x padVal [1] [1] [0]+                    y <- pad x padVal (v1 1) (v1 1) (v1 0)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [1.0, 2.0]@@ -210,7 +210,7 @@                 $ do                     x <- arg @'[4] @'F32                     idx <- constant @'[] @'I64 1-                    y <- dynamicSlice x [idx] [2]+                    y <- dynamicSlice x [idx] (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0]

@@ -5,6 +5,7 @@ module Test.Runtime.EndToEndMatmul where  import qualified Data.Vector.Storable as V+import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -83,7 +84,7 @@                 $ do                     x <- arg @'[1, 2, 3] @'F32                     y <- arg @'[3, 2] @'F32-                    z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y+                    z <- dotGeneral VS.empty VS.empty (v1 2) (v1 0) 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]

@@ -5,6 +5,7 @@ module Test.Runtime.EndToEndMatmulGPU (tests) where  import qualified Data.Vector.Storable as V+import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -84,7 +85,7 @@                 $ do                     x <- arg @'[1, 2, 3] @'F32                     y <- arg @'[3, 2] @'F32-                    z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y+                    z <- dotGeneral VS.empty VS.empty (v1 2) (v1 0) 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]

@@ -38,7 +38,7 @@                 [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                    y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) 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, 7.0, 8.0,@@ -54,7 +54,7 @@                 [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- avgPool [2, 2] [2, 2] x+                    y <- avgPool (v2 2 2) (v2 2 2) 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, 7.0, 8.0,

@@ -42,7 +42,7 @@                 [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                    y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) 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, 7.0, 8.0,@@ -57,7 +57,7 @@                 [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- avgPool [2, 2] [2, 2] x+                    y <- avgPool (v2 2 2) (v2 2 2) 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, 7.0, 8.0,

@@ -40,7 +40,7 @@                 [ FuncArg "arg0" (TensorType [2, 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [1, 0] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x                     return y         exec <- compile api client (render modu)         bufIn <- toDeviceF32 api client input2x2 [2, 2]@@ -52,7 +52,7 @@                 [ FuncArg "arg0" (TensorType [2, 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [0, 1] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x                     return y         exec <- compile api client (render modu)         bufIn <- toDeviceF32 api client input2x2 [2, 2]

@@ -44,7 +44,7 @@                 [ FuncArg "arg0" (TensorType [2, 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [1, 0] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x                     return y         exec <- compile api client (render modu)         bufIn <- toDeviceF32On api client dev input2x2 [2, 2]@@ -57,7 +57,7 @@                 [ FuncArg "arg0" (TensorType [2, 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [0, 1] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x                     return y         exec <- compile api client (render modu)         bufIn <- toDeviceF32On api client dev input2x2 [2, 2]