hbayesian (empty) → 0.1.0.0
raw patch · 37 files changed
+4339/−0 lines, 37 filesdep +basedep +bytestringdep +directory
Dependencies added: base, bytestring, directory, hbayesian, hhlo, tasty, tasty-golden, tasty-hunit, text, vector
Files
- CHANGELOG.md +5/−0
- LICENSE +21/−0
- README.md +274/−0
- app/Main.hs +35/−0
- examples/BivariateGaussianMALA.hs +201/−0
- examples/Common.hs +73/−0
- examples/CorrelatedGaussianHMC.hs +472/−0
- examples/CorrelatedGaussianHMCMain.hs +15/−0
- examples/GaussianProcessEllipticalSlice.hs +142/−0
- examples/LinearRegressionRandomWalk.hs +149/−0
- examples/LogisticRegressionHMC.hs +217/−0
- examples/Main.hs +101/−0
- hbayesian.cabal +178/−0
- src/HBayesian/Chain.hs +355/−0
- src/HBayesian/Core.hs +74/−0
- src/HBayesian/Diagnostics.hs +113/−0
- src/HBayesian/HHLO/Compile.hs +32/−0
- src/HBayesian/HHLO/Loops.hs +182/−0
- src/HBayesian/HHLO/Ops.hs +289/−0
- src/HBayesian/HHLO/PJRT.hs +60/−0
- src/HBayesian/HHLO/RNG.hs +96/−0
- src/HBayesian/InferenceLoop.hs +68/−0
- src/HBayesian/MCMC/EllipticalSlice.hs +58/−0
- src/HBayesian/MCMC/HMC.hs +113/−0
- src/HBayesian/MCMC/MALA.hs +31/−0
- src/HBayesian/MCMC/RandomWalk.hs +61/−0
- src/HBayesian/PPL.hs +183/−0
- test/Main.hs +25/−0
- test/Test/Chain.hs +33/−0
- test/Test/Core.hs +20/−0
- test/Test/CorrelatedGaussian.hs +213/−0
- test/Test/Examples.hs +47/−0
- test/Test/HHLO/Loops.hs +69/−0
- test/Test/HHLO/Ops.hs +70/−0
- test/Test/HHLO/RNG.hs +58/−0
- test/Test/MCMC.hs +150/−0
- test/Test/PPL.hs +56/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for hbayesian++## 0.1.0.0 -- YYYY-mm-dd++* First version. Released on an unsuspecting world.
+ LICENSE view
@@ -0,0 +1,21 @@+MIT License++Copyright (c) 2026 Le Niu++Permission is hereby granted, free of charge, to any person obtaining a copy+of this software and associated documentation files (the "Software"), to deal+in the Software without restriction, including without limitation the rights+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell+copies of the Software, and to permit persons to whom the Software is+furnished to do so, subject to the following conditions:++The above copyright notice and this permission notice shall be included in all+copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE+SOFTWARE.
+ README.md view
@@ -0,0 +1,274 @@+# HBayesian++Composable Bayesian inference in Haskell on StableHLO / XLA.++HBayesian provides MCMC samplers that compile to [StableHLO](https://github.com/openxla/stablehlo) and execute via PJRT (CPU, GPU, or TPU). It includes RandomWalk MH, Elliptical Slice, HMC, and MALA — along with chain combinators, diagnostics, and a shallow probabilistic programming layer.++## What makes HBayesian different++### Write Haskell, run on XLA++HBayesian is built on a simple idea: write ordinary Haskell functions that construct tensor computation graphs, then compile those graphs to StableHLO and execute them via PJRT.++You write ordinary Haskell functions that construct tensor computation graphs. Those graphs compile to self-contained StableHLO modules and execute on PJRT — the same runtime that powers JAX, TensorFlow, and PyTorch XLA. The result is that your sampler runs as a single XLA program, with the entire MCMC step (random number generation, proposal, gradient evaluation, acceptance) happening on device.++```haskell+-- This is not interpreted. It builds a StableHLO graph.+myLogPdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+myLogPdf theta = do+ alpha <- tslice1 @2 @'F32 theta 0+ beta <- tslice1 @2 @'F32 theta 1+ -- ... graph construction ...+```++### Compile-time shape safety++Tensors carry their shape and dtype in the type:++```haskell+x :: Tensor '[3] 'F32 -- 1-D float vector of length 3+y :: Tensor '[] 'F32 -- scalar float+```++Mismatches are caught at compile time. You cannot pass a `[3]` parameter vector to a sampler expecting `[2]` — the error appears in GHC, not at runtime on the device. This eliminates an entire class of "shape mismatch" bugs that are common in frameworks with dynamic typing.++### One language, end to end++Model definition, sampler design, and execution control are all Haskell. There is no Python/C++ boundary to cross, no foreign-function interface to debug, and no "works in the interpreter but fails on GPU" surprises. If your code type-checks, it compiles to StableHLO. If it compiles to StableHLO, it runs on PJRT.++### Composable samplers++Every sampler exposes the same `Kernel` interface:++```haskell+data Kernel s d state info = Kernel+ { kernelInit :: Key -> Tensor s d -> Builder (state s d)+ , kernelStep :: Key -> state s d -> Builder (state s d, info s d)+ }+```++This means the host loop is agnostic to the algorithm inside. You can swap RandomWalk for HMC without changing your execution code. The library provides four kernels out of the box:++| Sampler | Needs gradient? | Best for |+|---------|-----------------|----------|+| **RandomWalk** MH | No | Low dimensions, simple models |+| **EllipticalSlice** | No | Gaussian priors, no tuning needed |+| **HMC** | Yes | High dimensions, informed proposals |+| **MALA** | Yes | Cheap gradient steps |++### Chain combinators++Sampling is configured through ordinary function composition:++```haskell+sampleChain ck [0.0, 0.0] $+ burnIn 500 $ thin 2 $ withSeed 42 $ defaultChainConfig+ { ccNumIterations = 2000 }+```++`parallelChains` runs multiple independent chains with distinct seeds — essential for Gelman-Rubin convergence diagnostics.++### Diagnostics++`HBayesian.Diagnostics` operates on the `Diagnostic` records collected by `sampleChain`:++- **Acceptance rate** — fraction of accepted proposals+- **Gelman-Rubin R-hat** — across parallel chains+- **Effective sample size (ESS)** — autocorrelation-adjusted++### PPL layer++Write models as generative stories instead of manual log-densities:++```haskell+myModel :: PPL 2 ()+myModel = do+ alpha <- param 0+ beta <- param 1+ observe "alpha_prior" (normal 0.0 1.0) alpha+ observe "beta_prior" (normal 0.0 1.0) beta+```++`runPPL` desugars the story into a standard log-posterior function.++## Quick start++### Install the PJRT plugin++```bash+./scripts/pjrt_script.sh+```++This downloads `libpjrt_cpu.so` into `deps/pjrt/`. To use a custom plugin:++```bash+export HBAYESIAN_PJRT_PLUGIN=/path/to/libpjrt_cpu.so+```++### Build++```bash+cabal build+```++### Run a demo++```bash+# Run all 4 basic examples+cabal run hbayesian-examples -- --execute++# Run the HMC goodness-of-fit regression test+cabal run correlated-gaussian-hmc+```++## Usage++### Basic sampling++```haskell+import HBayesian.Chain+import HBayesian.MCMC.RandomWalk+import HBayesian.Diagnostics (acceptanceRate)+import qualified LinearRegressionRandomWalk as Ex++main :: IO ()+main = do+ let kernel = randomWalk Ex.linearRegLogPdf (RWConfig 0.1)+ ck = compileSimpleKernel kernel Ex.linearRegLogPdf+ (samples, diags) <- sampleChain ck [0.0, 0.0] $+ burnIn 100 $ thin 2 $ defaultChainConfig { ccNumIterations = 1000 }+ print (head samples)+ putStrLn $ "Acceptance rate: " ++ show (acceptanceRate diags)+```++### HMC with gradients++```haskell+import HBayesian.Chain+import HBayesian.MCMC.HMC++main :: IO ()+main = do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 10 }+ kernel = hmc myLogPdf myGradient config+ ck = compileHMC kernel myLogPdf myGradient+ (samples, diags) <- sampleChain ck (replicate 5 0.0) $+ burnIn 500 $ defaultChainConfig { ccNumIterations = 2000 }+ print (head samples)+```++### Parallel chains++```haskell+results <- parallelChains 4 (map (+ 0.5)) ck pos0 config+let chains = map fst results+```++### PPL layer++```haskell+import HBayesian.PPL++myModel :: PPL 2 ()+myModel = do+ alpha <- param 0+ beta <- param 1+ observe "alpha_prior" (normal 0.0 1.0) alpha+ observe "beta_prior" (normal 0.0 1.0) beta++logpdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+logpdf = runPPL myModel+```++## Examples++| Example | Sampler | What it shows | Run |+|---------|---------|---------------|-----|+| `LinearRegressionRandomWalk` | RandomWalk MH | Bayesian linear regression, 2-D params | `cabal run hbayesian-examples -- --execute` |+| `GaussianProcessEllipticalSlice` | Elliptical Slice | GP regression with Gaussian prior | `cabal run hbayesian-examples -- --execute` |+| `LogisticRegressionHMC` | HMC | Logistic regression with user-provided gradient | `cabal run hbayesian-examples -- --execute` |+| `BivariateGaussianMALA` | MALA | 2-D correlated Gaussian target | `cabal run hbayesian-examples -- --execute` |+| `CorrelatedGaussianHMC` | HMC | 5-D Gaussian with statistical GoF validation | `cabal run correlated-gaussian-hmc` |++Each example exposes `makeKernel` and `runChain`. Import them in GHCi to experiment interactively:++```bash+cabal repl hbayesian-examples+```++```haskell+import qualified LinearRegressionRandomWalk as Ex+Ex.runChain+```++### CorrelatedGaussianHMC — GoF regression test++`CorrelatedGaussianHMC` is the most rigorous example in the suite. It targets a 5-D correlated Gaussian with AR(1) covariance (ρ = 0.7) where the precision matrix and gradient are derived analytically. Because the model is exact, any deviation in the samples points directly to a sampler bug.++Running it prints a statistical report with four checks:++- **Marginal moments** — each dimension's mean and variance match N(μᵢ, 1)+- **Kolmogorov–Smirnov** — each marginal passes a KS test against the theoretical CDF+- **Mahalanobis χ²** — mean of (x−μ)ᵀΛ(x−μ) ≈ 5+- **Gelman–Rubin R-hat** — across 4 parallel chains, R-hat < 1.1 for every dimension++```bash+cabal run correlated-gaussian-hmc+```++Output:+```+==================================+ Goodness-of-Fit Report+==================================+Sample count: 2000++Marginal means (expected vs observed):+ dim 0: 1.0 vs 1.025 (diff: 2.54e-2) PASS+...++Mahalanobis distances:+ expected mean: 5.00+ observed mean: 5.04 PASS+```++## Tests++```bash+cabal test+```++34 tests covering Core, HHLO primitives, RNG, loops, all 4 samplers, chain combinators, PPL, and the CorrelatedGaussian GoF suite.++## Tutorial++A comprehensive two-level tutorial lives in `tutorial/TUTORIAL.md`:++- **Level 1** — Using inference algorithms (for end users)+- **Level 2** — Designing inference algorithms (for algorithm authors)++## Project structure++```+hbayesian/+|-- src/HBayesian/+| |-- Chain.hs -- Chain combinators+| |-- Diagnostics.hs -- MCMC diagnostics+| |-- PPL.hs -- Probabilistic programming layer+| |-- MCMC/ -- Samplers (RandomWalk, EllipticalSlice, HMC, MALA)+| |-- HHLO/ -- StableHLO primitives, RNG, loops, PJRT helpers+|-- examples/ -- Example suite + CLI entry points+|-- test/ -- Test suite+|-- tutorial/TUTORIAL.md+```++## Status++- **Phase 1** (Foundation): ✅ Core abstractions, RNG, loops, compilation+- **Phase 2** (Samplers): ✅ 4 samplers + practical examples+- **Phase 3** (Usability): ✅ Chain combinators, PPL, diagnostics, GoF validation+- **Phase 4+**: NUTS, adaptation, auto-diff — see `doc/` for design documents++## License++MIT
+ app/Main.hs view
@@ -0,0 +1,35 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Main (main) where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HBayesian.HHLO.Ops+import HBayesian.HHLO.RNG++main :: IO ()+main = do+ putStrLn "=== HBayesian Phase 1 Demo ==="++ -- Demo 1: Render a simple uniform RNG program+ let uniformModule = render $ moduleFromBuilder @'[4] @'F32 "uniform_demo"+ [ FuncArg "key" (TensorType [2] UI64) ] $ do+ k <- arg @'[2] @'UI64+ rngUniformF32 (Key k)+ putStrLn "\n-- Uniform RNG MLIR --"+ putStrLn (show uniformModule)++ -- Demo 2: Render a normal RNG program+ let normalModule = render $ moduleFromBuilder @'[4] @'F32 "normal_demo"+ [ FuncArg "key" (TensorType [2] UI64) ] $ do+ k <- arg @'[2] @'UI64+ rngNormalF32 (Key k)+ putStrLn "\n-- Normal RNG MLIR --"+ putStrLn (show normalModule)++ putStrLn "\n=== Demo complete ==="
+ examples/BivariateGaussianMALA.hs view
@@ -0,0 +1,201 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 4: Bivariate Gaussian Target with MALA.+--+-- A simple 2D correlated Gaussian where MALA's single leapfrog step+-- is competitive. The gradient is trivial and closed-form.+module BivariateGaussianMALA+ ( bivariateLogPdf+ , bivariateGrad+ , makeKernel+ , renderStepMlir+ , runChain+ , runChainV2+ ) where++import Data.Word (Word64)+import Data.Text (Text)+import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.HMC (HMCState(..))+import HBayesian.MCMC.MALA+import HBayesian.Chain+import Common++-- | Precision matrix Lambda = Sigma^{-1} for+-- Sigma = [[1.0, 0.8], [0.8, 1.0]]+lambda11, lambda12, lambda22 :: Float+lambda11 = 2.7778+lambda12 = -2.2222+lambda22 = 2.7778++mu1, mu2 :: Float+mu1 = 1.0+mu2 = 2.0++-- | Log-density of the bivariate Gaussian.+bivariateLogPdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+bivariateLogPdf theta = do+ mu <- buildMu+ diff <- tsub theta mu++ d1 <- tslice1 @2 @'F32 diff 0+ d2 <- tslice1 @2 @'F32 diff 1++ l11 <- tconstant @'[] @'F32 (realToFrac lambda11)+ l12 <- tconstant @'[] @'F32 (realToFrac lambda12)+ l22 <- tconstant @'[] @'F32 (realToFrac lambda22)++ term1a <- tmul l11 d1+ term1b <- tmul l12 d2+ term1 <- tadd term1a term1b+ quad1 <- tmul d1 term1++ term2a <- tmul l12 d1+ term2b <- tmul l22 d2+ term2 <- tadd term2a term2b+ quad2 <- tmul d2 term2++ quadForm <- tadd quad1 quad2+ negHalf <- tconstant @'[] @'F32 (-0.5)+ tmul negHalf quadForm++-- | Gradient of the bivariate Gaussian log-density.+-- grad = -Lambda * (theta - mu)+bivariateGrad :: Gradient '[2] 'F32+bivariateGrad theta = do+ mu <- buildMu+ diff <- tsub theta mu++ d1 <- tslice1 @2 @'F32 diff 0+ d2 <- tslice1 @2 @'F32 diff 1++ l11 <- tconstant @'[] @'F32 (realToFrac lambda11)+ l12 <- tconstant @'[] @'F32 (realToFrac lambda12)+ l22 <- tconstant @'[] @'F32 (realToFrac lambda22)++ g1a <- tmul l11 d1+ g1b <- tmul l12 d2+ g1Inner <- tadd g1a g1b+ g1 <- tnegate g1Inner++ g2a <- tmul l12 d1+ g2b <- tmul l22 d2+ g2Inner <- tadd g2a g2b+ g2 <- tnegate g2Inner++ tpack2 g1 g2++-- | Helper: build the mean vector as a constant tensor.+buildMu :: Builder (Tensor '[2] 'F32)+buildMu = do+ m1 <- tconstant @'[] @'F32 (realToFrac mu1)+ m2 <- tconstant @'[] @'F32 (realToFrac mu2)+ tpack2 m1 m2++-- | Factory: build a MALA kernel for this target.+makeKernel :: MALAConfig -> Kernel '[2] 'F32 (HMCState '[2] 'F32) (Info '[2] 'F32)+makeKernel config = mala bivariateLogPdf bivariateGrad config++-- | Tier A: render one kernel step to MLIR text.+renderStepMlir :: Text+renderStepMlir =+ renderKernelStep @'[2] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [2] F32)+ , FuncArg "p" (TensorType [2] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [2] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[2] @'F32+ p <- arg @'[2] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[2] @'F32+ let config = MALAConfig { malaStepSize = 0.1 }+ (state', _info) <- kernelStep (makeKernel config) (Key key) (HMCState pos p ld g)+ return (hmcPosition state')++-- | Tier B: run a short chain on PJRT and return sampled positions.+runChain :: IO [[Float]]+runChain = withPJRTCPU $ \api client -> do+ let config = MALAConfig { malaStepSize = 0.1 }+ kernel = makeKernel config++ -- Compile the log-pdf module+ let ldMod = moduleFromBuilder @'[] @'F32 "main"+ [ FuncArg "theta" (TensorType [2] F32) ] $ do+ theta <- arg @'[2] @'F32+ bivariateLogPdf theta+ ldExe <- compileModule api client ldMod++ -- Compile the gradient module+ let gradMod = moduleFromBuilder @'[2] @'F32 "main"+ [ FuncArg "theta" (TensorType [2] F32) ] $ do+ theta <- arg @'[2] @'F32+ bivariateGrad theta+ gradExe <- compileModule api client gradMod++ -- Compile the MALA step module (single result: position)+ let stepMod = moduleFromBuilder @'[2] @'F32 "main"+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [2] F32)+ , FuncArg "p" (TensorType [2] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [2] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[2] @'F32+ p <- arg @'[2] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[2] @'F32+ (state', _info) <- kernelStep kernel (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ stepExe <- compileModule api client stepMod++ let seed :: Word64 = 42+ theta0 = [0.0, 0.0]++ -- Compute initial log-density and gradient+ thetaBuf0 <- bufferFromF32 api client [2] theta0+ [ldBuf0] <- executeModule api ldExe [thetaBuf0]+ [ld0] <- bufferToF32 api ldBuf0 1+ [gBuf0] <- executeModule api gradExe [thetaBuf0]+ g0 <- bufferToF32 api gBuf0 2++ loop api client stepExe ldExe gradExe seed (0 :: Int) theta0 ld0 g0 (10 :: Int) []+ where+ loop _ _ _ _ _ _ _ _ _ _ 0 acc = return (reverse acc)+ loop api client stepExe ldExe gradExe seed step pos ld g n acc = do+ let key = [seed, fromIntegral step]+ zeroP = [0.0, 0.0]+ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client [2] pos+ pBuf <- bufferFromF32 api client [2] zeroP+ ldBuf <- bufferFromF32 api client [] [ld]+ gBuf <- bufferFromF32 api client [2] g+ [newPosBuf] <- executeModule api stepExe [keyBuf, posBuf, pBuf, ldBuf, gBuf]+ newPos <- bufferToF32 api newPosBuf 2+ -- Recompute log-density and gradient for the next step+ [newLdBuf] <- executeModule api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ [newGBuf] <- executeModule api gradExe [newPosBuf]+ newG <- bufferToF32 api newGBuf 2+ loop api client stepExe ldExe gradExe seed (step + 1) newPos newLd newG (n - 1) (newPos : acc)++-- | v0.2: Run a chain using the 'Chain' combinators.+runChainV2 :: IO ([[Float]], [Diagnostic])+runChainV2 = do+ let config = MALAConfig { malaStepSize = 0.1 }+ kernel = makeKernel config+ ck = compileHMC kernel bivariateLogPdf bivariateGrad+ sampleChain ck [0.0, 0.0] $ defaultChainConfig+ { ccNumIterations = 10+ , ccSeed = 42+ }
+ examples/Common.hs view
@@ -0,0 +1,73 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Shared utilities for the Phase 2 example suite.+module Common+ ( -- * Tier A: MLIR rendering+ renderKernelStep+ -- * Tier B: PJRT execution helpers+ , compileModule+ , executeModule+ , bufferFromF32+ , bufferFromUI64+ , bufferToF32+ , bufferToUI64+ ) where++import Data.Int (Int64)+import Data.Text (Text)+import Data.Vector.Storable (Vector)+import qualified Data.Vector.Storable as V+import Data.Word (Word64)+import Foreign.C.Types (CInt)++import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..), Module)+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import qualified Data.Text as T+import HHLO.Runtime.Buffer (toDevice, toDeviceF32, fromDevice, fromDeviceF32)+import HHLO.Runtime.Compile (compileWithOptions, defaultCompileOptions)+import HHLO.Runtime.Execute (execute)+import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTClient, PJRTExecutable, PJRTBuffer,+ bufferTypeU64)++-----------------------------------------------------------------------------+-- Tier A: MLIR rendering+-----------------------------------------------------------------------------++renderKernelStep :: forall s d. (KnownShape s, KnownDType d)+ => [FuncArg] -> Builder (Tensor s d) -> Text+renderKernelStep args b = render $ moduleFromBuilder @s @d "main" args b++-----------------------------------------------------------------------------+-- Tier B: PJRT execution+-----------------------------------------------------------------------------++-- | Compile a StableHLO 'Module' to a PJRT executable.+compileModule :: PJRTApi -> PJRTClient -> Module -> IO PJRTExecutable+compileModule api client modl =+ compileWithOptions api client (render modl) defaultCompileOptions++-- | Execute a compiled PJRT executable with input buffers.+executeModule :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer]+executeModule = execute++-- | Create a device buffer from a list of 'Float' values.+bufferFromF32 :: PJRTApi -> PJRTClient -> [Int] -> [Float] -> IO PJRTBuffer+bufferFromF32 api client dims vals =+ toDeviceF32 api client (V.fromList vals) (map fromIntegral dims)++-- | Create a device buffer from a list of 'Word64' values.+bufferFromUI64 :: PJRTApi -> PJRTClient -> [Int] -> [Word64] -> IO PJRTBuffer+bufferFromUI64 api client dims vals =+ toDevice api client (V.fromList vals) (map fromIntegral dims) bufferTypeU64++-- | Read back 'Float' values from a device buffer.+bufferToF32 :: PJRTApi -> PJRTBuffer -> Int -> IO [Float]+bufferToF32 api buf n = V.toList <$> fromDeviceF32 api buf n++-- | Read back 'Word64' values from a device buffer.+bufferToUI64 :: PJRTApi -> PJRTBuffer -> Int -> IO [Word64]+bufferToUI64 api buf n = V.toList <$> fromDevice api buf n
@@ -0,0 +1,472 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example: HMC on a 5-D correlated Gaussian with AR(1) covariance.+--+-- This example serves as a rigorous regression test for the HMC+-- implementation. The target distribution is known analytically, so+-- we can apply statistical goodness-of-fit tests to the samples.+--+-- Target: θ ~ N(μ, Σ) where Σ_ij = ρ^|i-j| (AR(1), uniform variance)+--+-- The precision matrix Λ = Σ⁻¹ is tridiagonal, making both log-density+-- and gradient easy to write in Builder.+module CorrelatedGaussianHMC+ ( -- * Ground truth+ targetMean+ , targetRho+ , targetDim+ , targetCov+ , targetPrecision+ -- * Model+ , gaussianLogPdf+ , gaussianGrad+ , makeKernel+ -- * HMC config (re-exported for test use)+ , HMCConfig (..)+ -- * Execution+ , runChain+ , runChainV2+ -- * Goodness-of-fit+ , goodnessOfFitReport+ -- * Debug+ , renderStepMlir+ ) where++import Data.List (sort)+import Data.Text (Text)+import Data.Word (Word64)+import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops hiding (map, maximum, sort)+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.HMC+import HBayesian.Chain+import Common++-----------------------------------------------------------------------------+-- Ground truth+-----------------------------------------------------------------------------++-- | Dimensionality.+targetDim :: Int+targetDim = 5++-- | True mean vector μ.+targetMean :: [Float]+targetMean = [1.0, -0.5, 2.0, 0.0, -1.0]++-- | AR(1) correlation coefficient.+targetRho :: Float+targetRho = 0.7++-- | Convenience: ρ².+rho2 :: Float+rho2 = targetRho * targetRho++-- | Precision matrix entries (analytical form for AR(1) with uniform variance).+-- Λ_ii for boundaries (i = 0, 4).+precBoundary :: Float+precBoundary = 1.0 / (1.0 - rho2)++-- | Λ_ii for interior points (i = 1, 2, 3).+precInterior :: Float+precInterior = (1.0 + rho2) / (1.0 - rho2)++-- | Λ_i,i+1 = Λ_i+1,i (off-diagonal).+precOffDiag :: Float+precOffDiag = -targetRho / (1.0 - rho2)++-- | Covariance matrix Σ_ij = ρ^|i-j| (for reference / validation).+targetCov :: [[Float]]+targetCov =+ [ [ targetRho ^ abs (i - j) | j <- [0 .. targetDim - 1] ]+ | i <- [0 .. targetDim - 1] ]++-- | Precision matrix Λ = Σ⁻¹ (tridiagonal, analytical form).+targetPrecision :: [[Float]]+targetPrecision =+ [ [ case abs (i - j) of+ 0 | i == 0 || i == targetDim - 1 -> precBoundary+ | otherwise -> precInterior+ 1 -> precOffDiag+ _ -> 0.0+ | j <- [0 .. targetDim - 1] ]+ | i <- [0 .. targetDim - 1] ]++-----------------------------------------------------------------------------+-- Model: log-density+-----------------------------------------------------------------------------++-- | Compute the quadratic form (θ - μ)^T Λ (θ - μ) for tridiagonal Λ.+--+-- Since Λ is symmetric tridiagonal, this expands to:+-- Σ_i Λ_ii * d_i² + 2 * Σ_i Λ_i,i+1 * d_i * d_{i+1}+quadraticForm :: Tensor '[5] 'F32 -> Builder (Tensor '[] 'F32)+quadraticForm theta = do+ -- diffs d_i = theta_i - mu_i+ d0 <- diff theta 0+ d1 <- diff theta 1+ d2 <- diff theta 2+ d3 <- diff theta 3+ d4 <- diff theta 4++ -- Diagonal terms: Λ_ii * d_i²+ q0 <- diagTerm precBoundary d0+ q1 <- diagTerm precInterior d1+ q2 <- diagTerm precInterior d2+ q3 <- diagTerm precInterior d3+ q4 <- diagTerm precBoundary d4++ -- Off-diagonal terms: 2 * Λ_i,i+1 * d_i * d_{i+1}+ q01 <- offDiagTerm precOffDiag d0 d1+ q12 <- offDiagTerm precOffDiag d1 d2+ q23 <- offDiagTerm precOffDiag d2 d3+ q34 <- offDiagTerm precOffDiag d3 d4++ -- Sum everything (all terms are scalars)+ s1 <- tadd q0 q01+ s2 <- tadd s1 q1+ s3 <- tadd s2 q12+ s4 <- tadd s3 q2+ s5 <- tadd s4 q23+ s6 <- tadd s5 q3+ s7 <- tadd s6 q34+ tadd s7 q4+ where+ diff th idx = do+ thi <- tslice1 @5 @'F32 th (fromIntegral idx)+ mui <- tconstant (realToFrac (targetMean !! idx))+ tsub thi mui+ diagTerm lam di = do+ lamT <- tconstant (realToFrac lam)+ di2 <- tmul di di+ tmul lamT di2+ offDiagTerm lam di dj = do+ lamT <- tconstant (realToFrac lam)+ two <- tconstant 2.0+ prod <- tmul di dj+ tmp <- tmul lamT prod+ tmul two tmp++gaussianLogPdf :: Tensor '[5] 'F32 -> Builder (Tensor '[] 'F32)+gaussianLogPdf theta = do+ quad <- quadraticForm theta+ negHalf <- tconstant (-0.5)+ tmul negHalf quad++-----------------------------------------------------------------------------+-- Gradient+-----------------------------------------------------------------------------++-- | ∇_θ log p(θ) = -Λ (θ - μ)+--+-- For tridiagonal Λ:+-- g_0 = -(Λ_00*d0 + Λ_01*d1)+-- g_i = -(Λ_{i,i-1}*d_{i-1} + Λ_ii*d_i + Λ_{i,i+1}*d_{i+1})+-- g_4 = -(Λ_43*d3 + Λ_44*d4)+gaussianGrad :: Gradient '[5] 'F32+gaussianGrad theta = do+ d0 <- diff theta 0+ d1 <- diff theta 1+ d2 <- diff theta 2+ d3 <- diff theta 3+ d4 <- diff theta 4++ g0 <- grad0 d0 d1+ g1 <- grad1 d0 d1 d2+ g2 <- grad2 d1 d2 d3+ g3 <- grad3 d2 d3 d4+ g4 <- grad4 d3 d4++ g0r <- treshape @'[] @'[1] g0+ g1r <- treshape @'[] @'[1] g1+ g2r <- treshape @'[] @'[1] g2+ g3r <- treshape @'[] @'[1] g3+ g4r <- treshape @'[] @'[1] g4+ concatenate @'[1] @'[5] @'F32 0 [g0r, g1r, g2r, g3r, g4r]+ where+ diff th idx = do+ thi <- tslice1 @5 @'F32 th (fromIntegral idx)+ mui <- tconstant (realToFrac (targetMean !! idx))+ tsub thi mui++ grad0 d0 d1 = do+ lam0 <- tconstant (realToFrac precBoundary)+ lam1 <- tconstant (realToFrac precOffDiag)+ t0 <- tmul lam0 d0+ t1 <- tmul lam1 d1+ s <- tadd t0 t1+ tnegate s++ grad1 d0 d1 d2 = do+ lam0 <- tconstant (realToFrac precOffDiag)+ lam1 <- tconstant (realToFrac precInterior)+ lam2 <- tconstant (realToFrac precOffDiag)+ t0 <- tmul lam0 d0+ t1 <- tmul lam1 d1+ t2 <- tmul lam2 d2+ s1 <- tadd t0 t1+ s <- tadd s1 t2+ tnegate s++ grad2 d1 d2 d3 = do+ lam1 <- tconstant (realToFrac precOffDiag)+ lam2 <- tconstant (realToFrac precInterior)+ lam3 <- tconstant (realToFrac precOffDiag)+ t1 <- tmul lam1 d1+ t2 <- tmul lam2 d2+ t3 <- tmul lam3 d3+ s1 <- tadd t1 t2+ s <- tadd s1 t3+ tnegate s++ grad3 d2 d3 d4 = do+ lam2 <- tconstant (realToFrac precOffDiag)+ lam3 <- tconstant (realToFrac precInterior)+ lam4 <- tconstant (realToFrac precOffDiag)+ t2 <- tmul lam2 d2+ t3 <- tmul lam3 d3+ t4 <- tmul lam4 d4+ s1 <- tadd t2 t3+ s <- tadd s1 t4+ tnegate s++ grad4 d3 d4 = do+ lam3 <- tconstant (realToFrac precOffDiag)+ lam4 <- tconstant (realToFrac precBoundary)+ t3 <- tmul lam3 d3+ t4 <- tmul lam4 d4+ s <- tadd t3 t4+ tnegate s++-----------------------------------------------------------------------------+-- Kernel+-----------------------------------------------------------------------------++makeKernel :: HMCConfig -> Kernel '[5] 'F32 (HMCState '[5] 'F32) (Info '[5] 'F32)+makeKernel config = hmc gaussianLogPdf gaussianGrad config++-----------------------------------------------------------------------------+-- Tier A: render MLIR+-----------------------------------------------------------------------------++renderStepMlir :: Text+renderStepMlir =+ renderKernelStep @'[5] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [5] F32)+ , FuncArg "p" (TensorType [5] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [5] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[5] @'F32+ p <- arg @'[5] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[5] @'F32+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 10 }+ (state', _info) <- kernelStep (makeKernel config) (Key key) (HMCState pos p ld g)+ return (hmcPosition state')++-----------------------------------------------------------------------------+-- Tier B: run chain (v0.1 style)+-----------------------------------------------------------------------------++runChain :: IO [[Float]]+runChain = withPJRTCPU $ \api client -> do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 10 }+ kernel = makeKernel config++ let ldMod = moduleFromBuilder @'[] @'F32 "main"+ [ FuncArg "theta" (TensorType [5] F32) ] $ do+ theta <- arg @'[5] @'F32+ gaussianLogPdf theta+ ldExe <- compileModule api client ldMod++ let gradMod = moduleFromBuilder @'[5] @'F32 "main"+ [ FuncArg "theta" (TensorType [5] F32) ] $ do+ theta <- arg @'[5] @'F32+ gaussianGrad theta+ gradExe <- compileModule api client gradMod++ let stepMod = moduleFromBuilder @'[5] @'F32 "main"+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [5] F32)+ , FuncArg "p" (TensorType [5] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [5] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[5] @'F32+ p <- arg @'[5] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[5] @'F32+ (state', _info) <- kernelStep kernel (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ stepExe <- compileModule api client stepMod++ let seed :: Word64 = 42+ pos0 = replicate targetDim 0.0++ posBuf0 <- bufferFromF32 api client [targetDim] pos0+ [ldBuf0] <- executeModule api ldExe [posBuf0]+ [ld0] <- bufferToF32 api ldBuf0 1+ [gBuf0] <- executeModule api gradExe [posBuf0]+ g0 <- bufferToF32 api gBuf0 targetDim++ loop api client stepExe ldExe gradExe seed (0 :: Int) pos0 ld0 g0 (10 :: Int) []+ where+ loop _ _ _ _ _ _ _ _ _ _ 0 acc = return (reverse acc)+ loop api client stepExe ldExe gradExe seed step pos ld g n acc = do+ let key = [seed, fromIntegral step]+ zeroP = replicate targetDim 0.0+ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client [targetDim] pos+ pBuf <- bufferFromF32 api client [targetDim] zeroP+ ldBuf <- bufferFromF32 api client [] [ld]+ gBuf <- bufferFromF32 api client [targetDim] g+ [newPosBuf] <- executeModule api stepExe [keyBuf, posBuf, pBuf, ldBuf, gBuf]+ newPos <- bufferToF32 api newPosBuf targetDim+ [newLdBuf] <- executeModule api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ [newGBuf] <- executeModule api gradExe [newPosBuf]+ newG <- bufferToF32 api newGBuf targetDim+ loop api client stepExe ldExe gradExe seed (step + 1) newPos newLd newG (n - 1) (newPos : acc)++-----------------------------------------------------------------------------+-- Tier B: run chain (v0.2 style)+-----------------------------------------------------------------------------++runChainV2 :: IO ([[Float]], [Diagnostic])+runChainV2 = do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 10 }+ kernel = makeKernel config+ ck = compileHMC kernel gaussianLogPdf gaussianGrad+ sampleChain ck (replicate targetDim 0.0) $+ burnIn 500 $ thin 2 $ defaultChainConfig+ { ccNumIterations = 2000+ , ccSeed = 42+ }++-----------------------------------------------------------------------------+-- Goodness-of-fit report+-----------------------------------------------------------------------------++meanF :: [Float] -> Float+meanF xs = sum xs / fromIntegral (length xs)++varianceF :: [Float] -> Float+varianceF xs =+ let m = meanF xs+ n = fromIntegral (length xs)+ in if n <= 1 then 0.0+ else sum [ (x - m) * (x - m) | x <- xs ] / (n - 1)++-- | Normal CDF approximation (Abramowitz & Stegun, formula 26.2.17).+normalCdf :: Float -> Float+normalCdf x =+ let z = realToFrac x :: Double+ b1 = 0.319381530+ b2 = -0.356563782+ b3 = 1.781477937+ b4 = -1.821255978+ b5 = 1.330274429+ p = 0.2316419+ c = 0.39894228+ t = 1.0 / (1.0 + p * abs z)+ phi = c * exp (-0.5 * z * z)+ poly = t * (b1 + t * (b2 + t * (b3 + t * (b4 + t * b5))))+ y = 1.0 - phi * poly+ in realToFrac $ if z > 0 then y else 1.0 - y++empiricalCdf :: Float -> [Float] -> Float+empiricalCdf x sortedXs =+ let n = fromIntegral (length sortedXs)+ count = fromIntegral (length (takeWhile (<= x) sortedXs))+ in count / n++ksStatistic :: (Float -> Float) -> [Float] -> Float+ksStatistic cdf xs =+ let sorted = sort xs+ n = fromIntegral (length sorted)+ empiricalCdfPrev x sortedXs =+ let count = fromIntegral (length (takeWhile (< x) sortedXs))+ in count / n+ diffs = [ max (abs (empiricalCdf x sorted - cdf x))+ (abs (empiricalCdfPrev x sorted - cdf x))+ | x <- sorted ]+ in maximum diffs++ksCriticalValue :: Int -> Float+ksCriticalValue n = 1.628 / sqrt (fromIntegral n)++mahalanobisSq :: [Float] -> [Float] -> [[Float]] -> Float+mahalanobisSq x mu lam =+ let diff = zipWith (-) x mu+ n = length diff+ in sum [ diff !! i * lam !! i !! j * diff !! j+ | i <- [0..n-1], j <- [0..n-1] ]++-- | Print a formatted goodness-of-fit report for HMC samples.+goodnessOfFitReport :: [[Float]] -> IO ()+goodnessOfFitReport samples = do+ let n = length samples+ putStrLn "=================================="+ putStrLn " Goodness-of-Fit Report"+ putStrLn "=================================="+ putStrLn $ "Sample count: " ++ show n+ putStrLn ""++ -- Marginal means+ putStrLn "Marginal means (expected vs observed):"+ forM_ (zip [0..4] targetMean) $ \(i, muTrue) -> do+ let xs = map (!! i) samples+ muHat = meanF xs+ err = abs (muHat - muTrue)+ se = 1.0 / sqrt (fromIntegral n)+ status = if err < 3.0 * se then "PASS" else "FAIL"+ putStrLn $ " dim " ++ show i ++ ": " ++ padL 6 (show muTrue)+ ++ " vs " ++ padL 6 (show muHat)+ ++ " (diff: " ++ padL 6 (show err) ++ ") " ++ status++ putStrLn ""+ putStrLn "Marginal variances (expected vs observed):"+ forM_ [0..4] $ \i -> do+ let xs = map (!! i) samples+ varHat = varianceF xs+ err = abs (varHat - 1.0)+ status = if err < 0.15 then "PASS" else "FAIL"+ putStrLn $ " dim " ++ show i ++ ": " ++ padL 6 ("1.0")+ ++ " vs " ++ padL 6 (show varHat)+ ++ " (diff: " ++ padL 6 (show err) ++ ") " ++ status++ putStrLn ""+ let crit = ksCriticalValue n+ putStrLn $ "KS tests (critical value: " ++ show crit ++ "):"+ forM_ [0..4] $ \i -> do+ let xs = map (!! i) samples+ mu = targetMean !! i+ cdf x = normalCdf ((x - mu) / 1.0)+ stat = ksStatistic cdf xs+ status = if stat < crit then "PASS" else "FAIL"+ putStrLn $ " dim " ++ show i ++ ": stat=" ++ padL 8 (show stat)+ ++ " " ++ status++ putStrLn ""+ let dists = [ mahalanobisSq s targetMean targetPrecision | s <- samples ]+ meanD = meanF dists+ se = sqrt (10.0 / fromIntegral n)+ thresh = 3.0 * se+ err = abs (meanD - 5.0)+ status = if err < thresh then "PASS" else "FAIL"+ putStrLn "Mahalanobis distances:"+ putStrLn $ " expected mean: 5.00"+ putStrLn $ " observed mean: " ++ show meanD ++ " " ++ status+ putStrLn ""+ where+ padL w s = replicate (max 0 (w - length s)) ' ' ++ s+ forM_ = flip mapM_
@@ -0,0 +1,15 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Main (main) where++import HBayesian.Diagnostics (acceptanceRate)+import qualified CorrelatedGaussianHMC as Ex++main :: IO ()+main = do+ putStrLn "Running HMC on 5-D correlated Gaussian (AR(1), rho=0.7)..."+ putStrLn ""+ (samples, diags) <- Ex.runChainV2+ Ex.goodnessOfFitReport samples+ putStrLn $ "Acceptance rate: " ++ show (acceptanceRate diags)
+ examples/GaussianProcessEllipticalSlice.hs view
@@ -0,0 +1,142 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 2: Gaussian Process Regression with Elliptical Slice Sampling.+--+-- We use a tiny n=3 dataset with an identity covariance prior+-- (K = I, so K^{-1} = I) to avoid needing a Cholesky decomposition+-- in the Builder. The example still demonstrates ESS on a Gaussian+-- prior, which is the sampler's natural habitat.+module GaussianProcessEllipticalSlice+ ( yData+ , gpLogPdf+ , makeKernel+ , renderStepMlir+ , runChain+ , runChainV2+ ) where++import Data.Word (Word64)+import Data.Text (Text)+import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.EllipticalSlice+import HBayesian.Chain+import Common++-- | Observed responses (n = 3).+yData :: [Float]+yData = [0.5, 2.0, 3.5]++-- | GP log-posterior with identity prior covariance.+gpLogPdf :: Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32)+gpLogPdf f = do+ y0 <- tconstant @'[] @'F32 (realToFrac (yData !! 0))+ y1 <- tconstant @'[] @'F32 (realToFrac (yData !! 1))+ y2 <- tconstant @'[] @'F32 (realToFrac (yData !! 2))++ f0 <- tslice1 @3 @'F32 f 0+ f1 <- tslice1 @3 @'F32 f 1+ f2 <- tslice1 @3 @'F32 f 2++ d0 <- tsub y0 f0+ d1 <- tsub y1 f1+ d2 <- tsub y2 f2++ d0sq <- tmul d0 d0+ d1sq <- tmul d1 d1+ d2sq <- tmul d2 d2++ negHalf <- tconstant @'[] @'F32 (-0.5)+ llh0 <- tmul negHalf d0sq+ llh1 <- tmul negHalf d1sq+ llh2 <- tmul negHalf d2sq++ llh01 <- tadd llh0 llh1+ llh <- tadd llh01 llh2++ f0sq <- tmul f0 f0+ f1sq <- tmul f1 f1+ f2sq <- tmul f2 f2++ fSqSum <- tadd f0sq =<< tadd f1sq f2sq+ prior <- tmul negHalf fSqSum++ tadd llh prior++-- | Factory: build an Elliptical Slice kernel for this model.+makeKernel :: SimpleKernel '[3] 'F32+makeKernel = ellipticalSlice gpLogPdf++-- | Tier A: render one kernel step to MLIR text.+renderStepMlir :: Text+renderStepMlir =+ renderKernelStep @'[3] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [3] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[3] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep makeKernel (Key key) (State pos ld)+ return (statePosition state')++-- | Tier B: run a short chain on PJRT and return the sampled latent vectors.+runChain :: IO [[Float]]+runChain = withPJRTCPU $ \api client -> do+ -- Compile the log-pdf module+ let ldMod = moduleFromBuilder @'[] @'F32 "main"+ [ FuncArg "f" (TensorType [3] F32) ] $ do+ f <- arg @'[3] @'F32+ gpLogPdf f+ ldExe <- compileModule api client ldMod++ -- Compile the kernel-step module (single result: position)+ let stepMod = moduleFromBuilder @'[3] @'F32 "main"+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [3] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[3] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep makeKernel (Key key) (State pos ld)+ return (statePosition state')+ stepExe <- compileModule api client stepMod++ let seed :: Word64 = 42+ f0 = [0.0, 0.0, 0.0]++ -- Compute initial log-density+ fBuf0 <- bufferFromF32 api client [3] f0+ [ldBuf0] <- executeModule api ldExe [fBuf0]+ [ld0] <- bufferToF32 api ldBuf0 1++ loop api client stepExe ldExe seed (0 :: Int) f0 ld0 (10 :: Int) []+ where+ loop _ _ _ _ _ _ _ _ 0 acc = return (reverse acc)+ loop api client stepExe ldExe seed step pos ld n acc = do+ let key = [seed, fromIntegral step]+ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client [3] pos+ ldBuf <- bufferFromF32 api client [] [ld]+ [newPosBuf] <- executeModule api stepExe [keyBuf, posBuf, ldBuf]+ newPos <- bufferToF32 api newPosBuf 3+ [newLdBuf] <- executeModule api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ loop api client stepExe ldExe seed (step + 1) newPos newLd (n - 1) (newPos : acc)++-- | v0.2: Run a chain using the 'Chain' combinators.+runChainV2 :: IO ([[Float]], [Diagnostic])+runChainV2 = do+ let ck = compileSimpleKernel makeKernel gpLogPdf+ sampleChain ck [0.0, 0.0, 0.0] $ defaultChainConfig+ { ccNumIterations = 10+ , ccSeed = 42+ }
+ examples/LinearRegressionRandomWalk.hs view
@@ -0,0 +1,149 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 1: Bayesian Linear Regression with RandomWalk MH.+--+-- Model: y_i = alpha + beta * x_i + epsilon_i, epsilon_i ~ N(0, 0.25)+-- Prior: alpha ~ N(0, 1), beta ~ N(0, 1)+module LinearRegressionRandomWalk+ ( dataset+ , linearRegLogPdf+ , makeKernel+ , renderStepMlir+ , runChain+ , runChainV2+ ) where++import Data.Word (Word64)+import Data.Text (Text)+import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.RandomWalk+import HBayesian.Chain+import Common++-- | Fixed synthetic dataset (n = 5).+dataset :: [(Float, Float)]+dataset =+ [ (0.0, 0.5)+ , (1.0, 2.0)+ , (2.0, 3.5)+ , (3.0, 5.0)+ , (4.0, 6.5)+ ]++-- | Log-posterior for Bayesian linear regression.+linearRegLogPdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+linearRegLogPdf theta = do+ alpha <- tslice1 @2 @'F32 theta 0+ beta <- tslice1 @2 @'F32 theta 1++ let likelihoodPoint (x, y) = do+ xT <- tconstant @'[] @'F32 (realToFrac x)+ yT <- tconstant @'[] @'F32 (realToFrac y)+ betaX <- tmul beta xT+ predVal <- tadd alpha betaX+ diff <- tsub yT predVal+ diffSq <- tmul diff diff+ negTwo <- tconstant @'[] @'F32 (-2.0)+ tmul negTwo diffSq++ llh0 <- likelihoodPoint (dataset !! 0)+ llh1 <- likelihoodPoint (dataset !! 1)+ llh2 <- likelihoodPoint (dataset !! 2)+ llh3 <- likelihoodPoint (dataset !! 3)+ llh4 <- likelihoodPoint (dataset !! 4)++ llh01 <- tadd llh0 llh1+ llh23 <- tadd llh2 llh3+ llh0123 <- tadd llh01 llh23+ llh <- tadd llh0123 llh4++ alphaSq <- tmul alpha alpha+ betaSq <- tmul beta beta+ negHalf <- tconstant @'[] @'F32 (-0.5)+ priorAlpha <- tmul negHalf alphaSq+ priorBeta <- tmul negHalf betaSq++ tadd llh =<< tadd priorAlpha priorBeta++-- | Factory: build a RandomWalk kernel for this model.+makeKernel :: RWConfig -> SimpleKernel '[2] 'F32+makeKernel config = randomWalk linearRegLogPdf config++-- | Tier A: render one kernel step to MLIR text.+renderStepMlir :: Text+renderStepMlir =+ renderKernelStep @'[2] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [2] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[2] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep (makeKernel (RWConfig 0.1)) (Key key) (State pos ld)+ return (statePosition state')++-- | Tier B: run a short chain on PJRT and return the sampled positions.+runChain :: IO [[Float]]+runChain = withPJRTCPU $ \api client -> do+ let kernel = makeKernel (RWConfig 0.1)++ -- Compile the log-pdf module+ let ldMod = moduleFromBuilder @'[] @'F32 "main"+ [ FuncArg "pos" (TensorType [2] F32) ] $ do+ pos <- arg @'[2] @'F32+ linearRegLogPdf pos+ ldExe <- compileModule api client ldMod++ -- Compile the kernel-step module (single result: position)+ let stepMod = moduleFromBuilder @'[2] @'F32 "main"+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [2] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[2] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep kernel (Key key) (State pos ld)+ return (statePosition state')+ stepExe <- compileModule api client stepMod++ let seed :: Word64 = 42+ pos0 = [0.0, 0.0]++ -- Compute initial log-density+ posBuf0 <- bufferFromF32 api client [2] pos0+ [ldBuf0] <- executeModule api ldExe [posBuf0]+ [ld0] <- bufferToF32 api ldBuf0 1++ loop api client stepExe ldExe seed (0 :: Int) pos0 ld0 (10 :: Int) []+ where+ loop _ _ _ _ _ _ _ _ 0 acc = return (reverse acc)+ loop api client stepExe ldExe seed step pos ld n acc = do+ let key = [seed, fromIntegral step]+ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client [2] pos+ ldBuf <- bufferFromF32 api client [] [ld]+ [newPosBuf] <- executeModule api stepExe [keyBuf, posBuf, ldBuf]+ newPos <- bufferToF32 api newPosBuf 2+ -- Recompute log-density for the next step+ [newLdBuf] <- executeModule api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ loop api client stepExe ldExe seed (step + 1) newPos newLd (n - 1) (newPos : acc)++-- | v0.2: Run a chain using the 'Chain' combinators.+runChainV2 :: IO ([[Float]], [Diagnostic])+runChainV2 = do+ let kernel = makeKernel (RWConfig 0.1)+ ck = compileSimpleKernel kernel linearRegLogPdf+ sampleChain ck [0.0, 0.0] $ defaultChainConfig+ { ccNumIterations = 10+ , ccSeed = 42+ }
+ examples/LogisticRegressionHMC.hs view
@@ -0,0 +1,217 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 3: Bayesian Logistic Regression with HMC.+--+-- Binary classification with D=3 features and n=4 observations.+-- The user provides the gradient explicitly, demonstrating the+-- manual-gradient path before Phase 5 auto-diff.+module LogisticRegressionHMC+ ( logisticRegLogPdf+ , logisticRegGrad+ , makeKernel+ , renderStepMlir+ , runChain+ , runChainV2+ ) where++import Data.Word (Word64)+import Data.Text (Text)+import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.HMC+import HBayesian.Chain+import Common++-- | Log-posterior for Bayesian logistic regression.+logisticRegLogPdf :: Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32)+logisticRegLogPdf beta = do+ let mkConstX [a, b, c] = do+ ca <- tconstant @'[] @'F32 (realToFrac a)+ cb <- tconstant @'[] @'F32 (realToFrac b)+ cc <- tconstant @'[] @'F32 (realToFrac c)+ tpack3 ca cb cc+ mkConstX _ = error "mkConstX: expected exactly 3 elements"++ x0 <- mkConstX [1.0 :: Float, 0.5, -0.5]+ y0 <- tconstant @'[] @'F32 1.0+ lp0 <- logLikPoint x0 y0 beta++ x1 <- mkConstX [1.0 :: Float, 1.0, -1.0]+ y1 <- tconstant @'[] @'F32 1.0+ lp1 <- logLikPoint x1 y1 beta++ x2 <- mkConstX [1.0 :: Float, 1.5, -1.5]+ y2 <- tconstant @'[] @'F32 0.0+ lp2 <- logLikPoint x2 y2 beta++ x3 <- mkConstX [1.0 :: Float, 2.0, -2.0]+ y3 <- tconstant @'[] @'F32 0.0+ lp3 <- logLikPoint x3 y3 beta++ llh01 <- tadd lp0 lp1+ llh23 <- tadd lp2 lp3+ llh <- tadd llh01 llh23++ betaSq <- tmul beta beta+ betaSumSq <- tsumAll betaSq+ negHalf <- tconstant @'[] @'F32 (-0.5)+ prior <- tmul negHalf betaSumSq++ tadd llh prior++-- | Single data-point log-likelihood contribution.+logLikPoint :: Tensor '[3] 'F32 -> Tensor '[] 'F32 -> Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32)+logLikPoint x_i y_i beta = do+ dotProd <- tsumAll =<< tmul x_i beta+ sig <- tsigmoid dotProd+ one <- tconstant @'[] @'F32 1.0+ logSig <- tlog sig+ logOneMinusSig <- tlog =<< tsub one sig+ term1 <- tmul y_i logSig+ oneMinusY <- tsub one y_i+ term2 <- tmul oneMinusY logOneMinusSig+ tadd term1 term2++-- | User-provided gradient of the log-posterior.+logisticRegGrad :: Gradient '[3] 'F32+logisticRegGrad beta = do+ let mkConstX [a, b, c] = do+ ca <- tconstant @'[] @'F32 (realToFrac a)+ cb <- tconstant @'[] @'F32 (realToFrac b)+ cc <- tconstant @'[] @'F32 (realToFrac c)+ tpack3 ca cb cc+ mkConstX _ = error "mkConstX: expected exactly 3 elements"++ grad0 <- gradPoint [1.0, 0.5, -0.5] 1.0 beta mkConstX+ grad1 <- gradPoint [1.0, 1.0, -1.0] 1.0 beta mkConstX+ grad2 <- gradPoint [1.0, 1.5, -1.5] 0.0 beta mkConstX+ grad3 <- gradPoint [1.0, 2.0, -2.0] 0.0 beta mkConstX++ g01 <- tadd grad0 grad1+ g23 <- tadd grad2 grad3+ gradLik <- tadd g01 g23++ negOne <- tconstant @'[3] @'F32 (-1.0)+ gradPrior <- tmul negOne beta++ tadd gradLik gradPrior++-- | Gradient contribution from a single data point.+gradPoint :: [Float] -> Float -> Tensor '[3] 'F32+ -> ([Float] -> Builder (Tensor '[3] 'F32))+ -> Builder (Tensor '[3] 'F32)+gradPoint xs yi beta mkConstX = do+ x_i <- mkConstX xs+ dotProd <- tsumAll =<< tmul x_i beta+ sig <- tsigmoid dotProd+ yT <- tconstant @'[] @'F32 (realToFrac yi)+ residual <- tsub yT sig+ residBC <- tbroadcast @'[] @'[3] [] residual+ tmul x_i residBC++-- | Factory: build an HMC kernel for this model.+makeKernel :: HMCConfig -> Kernel '[3] 'F32 (HMCState '[3] 'F32) (Info '[3] 'F32)+makeKernel config = hmc logisticRegLogPdf logisticRegGrad config++-- | Tier A: render one kernel step to MLIR text.+renderStepMlir :: Text+renderStepMlir =+ renderKernelStep @'[3] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [3] F32)+ , FuncArg "p" (TensorType [3] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [3] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[3] @'F32+ p <- arg @'[3] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[3] @'F32+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 2 }+ (state', _info) <- kernelStep (makeKernel config) (Key key) (HMCState pos p ld g)+ return (hmcPosition state')++-- | Tier B: run a short chain on PJRT and return sampled beta vectors.+runChain :: IO [[Float]]+runChain = withPJRTCPU $ \api client -> do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 2 }+ kernel = makeKernel config++ -- Compile the log-pdf module+ let ldMod = moduleFromBuilder @'[] @'F32 "main"+ [ FuncArg "beta" (TensorType [3] F32) ] $ do+ beta <- arg @'[3] @'F32+ logisticRegLogPdf beta+ ldExe <- compileModule api client ldMod++ -- Compile the gradient module+ let gradMod = moduleFromBuilder @'[3] @'F32 "main"+ [ FuncArg "beta" (TensorType [3] F32) ] $ do+ beta <- arg @'[3] @'F32+ logisticRegGrad beta+ gradExe <- compileModule api client gradMod++ -- Compile the HMC step module (single result: position)+ let stepMod = moduleFromBuilder @'[3] @'F32 "main"+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [3] F32)+ , FuncArg "p" (TensorType [3] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [3] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[3] @'F32+ p <- arg @'[3] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[3] @'F32+ (state', _info) <- kernelStep kernel (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ stepExe <- compileModule api client stepMod++ let seed :: Word64 = 42+ beta0 = [0.0, 0.0, 0.0]++ -- Compute initial log-density and gradient+ betaBuf0 <- bufferFromF32 api client [3] beta0+ [ldBuf0] <- executeModule api ldExe [betaBuf0]+ [ld0] <- bufferToF32 api ldBuf0 1+ [gBuf0] <- executeModule api gradExe [betaBuf0]+ g0 <- bufferToF32 api gBuf0 3++ loop api client stepExe ldExe gradExe seed (0 :: Int) beta0 ld0 g0 (10 :: Int) []+ where+ loop _ _ _ _ _ _ _ _ _ _ 0 acc = return (reverse acc)+ loop api client stepExe ldExe gradExe seed step pos ld g n acc = do+ let key = [seed, fromIntegral step]+ zeroP = [0.0, 0.0, 0.0]+ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client [3] pos+ pBuf <- bufferFromF32 api client [3] zeroP+ ldBuf <- bufferFromF32 api client [] [ld]+ gBuf <- bufferFromF32 api client [3] g+ [newPosBuf] <- executeModule api stepExe [keyBuf, posBuf, pBuf, ldBuf, gBuf]+ newPos <- bufferToF32 api newPosBuf 3+ -- Recompute log-density and gradient for the next step+ [newLdBuf] <- executeModule api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ [newGBuf] <- executeModule api gradExe [newPosBuf]+ newG <- bufferToF32 api newGBuf 3+ loop api client stepExe ldExe gradExe seed (step + 1) newPos newLd newG (n - 1) (newPos : acc)++-- | v0.2: Run a chain using the 'Chain' combinators.+runChainV2 :: IO ([[Float]], [Diagnostic])+runChainV2 = do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 2 }+ kernel = makeKernel config+ ck = compileHMC kernel logisticRegLogPdf logisticRegGrad+ sampleChain ck [0.0, 0.0, 0.0] $ defaultChainConfig+ { ccNumIterations = 10+ , ccSeed = 42+ }
+ examples/Main.hs view
@@ -0,0 +1,101 @@+module Main (main) where++import System.Environment (getArgs, setEnv)+import System.Exit (exitFailure)+import qualified Data.Text.IO as T+import qualified LinearRegressionRandomWalk as Ex1+import qualified GaussianProcessEllipticalSlice as Ex2+import qualified LogisticRegressionHMC as Ex3+import qualified BivariateGaussianMALA as Ex4++-----------------------------------------------------------------------------+-- CLI+-----------------------------------------------------------------------------++data Mode = ModeRender | ModeExecute+ deriving (Eq, Show)++data Options = Options+ { optMode :: !Mode+ , optPJRTPlugin :: !(Maybe FilePath)+ }++defaultOptions :: Options+defaultOptions = Options+ { optMode = ModeRender+ , optPJRTPlugin = Nothing+ }++usage :: String+usage = unlines+ [ "Usage: hbayesian-examples [OPTIONS]"+ , ""+ , "Options:"+ , " --render Print StableHLO MLIR for all examples (default)"+ , " --execute Run MCMC chains on PJRT"+ , " --pjrt-plugin PATH Use a custom PJRT plugin (default: deps/pjrt/libpjrt_cpu.so)"+ , " --help Show this message"+ ]++parseArgs :: [String] -> Either String Options+parseArgs = go defaultOptions+ where+ go opts [] = Right opts+ go opts ("--render" : rest) = go (opts { optMode = ModeRender }) rest+ go opts ("--execute" : rest) = go (opts { optMode = ModeExecute }) rest+ go opts ("--pjrt-plugin" : path : rest) = go (opts { optPJRTPlugin = Just path }) rest+ go _ ("--help" : _) = Left usage+ go _ (bad : _) = Left $ "Unknown flag: " ++ bad ++ "\n" ++ usage++-----------------------------------------------------------------------------+-- Main+-----------------------------------------------------------------------------++main :: IO ()+main = do+ args <- getArgs+ opts <- case parseArgs args of+ Left msg -> putStrLn msg >> exitFailure+ Right o -> return o++ case optMode opts of+ ModeRender -> runRender+ ModeExecute -> runExecute (optPJRTPlugin opts)++runRender :: IO ()+runRender = do+ putStrLn "=== Example 1: Bayesian Linear Regression (RandomWalk) ==="+ T.putStrLn Ex1.renderStepMlir++ putStrLn "\n=== Example 2: Gaussian Process (EllipticalSlice) ==="+ T.putStrLn Ex2.renderStepMlir++ putStrLn "\n=== Example 3: Logistic Regression (HMC) ==="+ T.putStrLn Ex3.renderStepMlir++ putStrLn "\n=== Example 4: Bivariate Gaussian (MALA) ==="+ T.putStrLn Ex4.renderStepMlir++runExecute :: Maybe FilePath -> IO ()+runExecute mPluginPath = do+ -- If a custom plugin path is given, expose it via the env var so that+ -- the example modules' withPJRTCPU picks it up.+ case mPluginPath of+ Just path -> setEnv "HBAYESIAN_PJRT_PLUGIN" path+ Nothing -> return ()++ putStrLn "=== Example 1: Bayesian Linear Regression (RandomWalk) ==="+ samples1 <- Ex1.runChain+ mapM_ print samples1++ putStrLn "\n=== Example 2: Gaussian Process (EllipticalSlice) ==="+ samples2 <- Ex2.runChain+ mapM_ print samples2++ putStrLn "\n=== Example 3: Logistic Regression (HMC) ==="+ samples3 <- Ex3.runChain+ mapM_ print samples3++ putStrLn "\n=== Example 4: Bivariate Gaussian (MALA) ==="+ samples4 <- Ex4.runChain+ mapM_ print samples4
+ hbayesian.cabal view
@@ -0,0 +1,178 @@+cabal-version: 3.0+name: hbayesian+version: 0.1.0.0+synopsis: Composable Bayesian inference in Haskell on StableHLO/XLA+description:+ HBayesian is a library of MCMC samplers that compiles inference kernels to+ StableHLO and executes them via PJRT (CPU, GPU, or TPU). It provides+ stateless, composable transition kernels for MCMC, built on top of the HHLO+ Haskell frontend. Features include RandomWalk MH, Elliptical Slice, HMC,+ MALA, chain combinators (burn-in, thinning, parallel chains), host-side+ diagnostics (R-hat, ESS), and a shallow probabilistic programming layer.+ .+ Every sampler is a pure Haskell function that builds a StableHLO graph via+ the HHLO EDSL. There is no interpreter fallback — if it compiles, it runs+ on device.+license: MIT+license-file: LICENSE+copyright: (c) 2025 hbayesian contributors+author: hbayesian contributors+maintainer: le.niu@hotmail.com+category: Math, Statistics, AI+build-type: Simple+tested-with: GHC >= 9.6.7++homepage: https://github.com/overshiki/hbayesian+bug-reports: https://github.com/overshiki/hbayesian/issues++extra-doc-files:+ CHANGELOG.md+ README.md++source-repository head+ type: git+ location: https://github.com/overshiki/hbayesian.git++common warnings+ ghc-options:+ -Wall+ -Wcompat+ -Widentities+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wmissing-export-lists+ -Wmissing-home-modules+ -Wpartial-fields+ -Wredundant-constraints++library+ import: warnings+ exposed-modules:+ HBayesian.Core+ HBayesian.HHLO.Ops+ HBayesian.HHLO.RNG+ HBayesian.HHLO.Loops+ HBayesian.HHLO.Compile+ HBayesian.HHLO.PJRT+ HBayesian.Chain+ HBayesian.Diagnostics+ HBayesian.PPL+ HBayesian.InferenceLoop+ HBayesian.MCMC.RandomWalk+ HBayesian.MCMC.EllipticalSlice+ HBayesian.MCMC.HMC+ HBayesian.MCMC.MALA+ build-depends:+ base >= 4.18.2.1 && < 4.20+ , directory >= 1.3.8.1 && < 1.4+ , hhlo >= 0.2.0.0 && < 0.3+ , text >= 2.0 && < 2.2+ , vector >= 0.13 && < 0.14+ hs-source-dirs: src+ default-language: Haskell2010+ default-extensions:+ DataKinds+ KindSignatures+ TypeApplications+ TypeFamilies+ TypeOperators+ ScopedTypeVariables++executable hbayesian-demo+ import: warnings+ main-is: Main.hs+ build-depends:+ base >= 4.18.2.1 && < 4.20+ , hbayesian+ , hhlo >= 0.2.0.0 && < 0.3+ , text >= 2.0 && < 2.2+ hs-source-dirs: app+ default-language: Haskell2010++executable hbayesian-examples+ import: warnings+ main-is: Main.hs+ other-modules:+ Common+ LinearRegressionRandomWalk+ GaussianProcessEllipticalSlice+ LogisticRegressionHMC+ BivariateGaussianMALA+ CorrelatedGaussianHMC+ build-depends:+ base >= 4.18.2.1 && < 4.20+ , hbayesian+ , hhlo >= 0.2.0.0 && < 0.3+ , text >= 2.0 && < 2.2+ , vector >= 0.13 && < 0.14+ hs-source-dirs: examples+ default-language: Haskell2010+ default-extensions:+ DataKinds+ KindSignatures+ TypeApplications+ TypeFamilies+ TypeOperators+ ScopedTypeVariables++executable correlated-gaussian-hmc+ import: warnings+ main-is: CorrelatedGaussianHMCMain.hs+ other-modules:+ Common+ CorrelatedGaussianHMC+ build-depends:+ base >= 4.18.2.1 && < 4.20+ , hbayesian+ , hhlo >= 0.2.0.0 && < 0.3+ , text >= 2.0 && < 2.2+ , vector >= 0.13 && < 0.14+ hs-source-dirs: examples+ default-language: Haskell2010+ default-extensions:+ DataKinds+ KindSignatures+ TypeApplications+ TypeFamilies+ TypeOperators+ ScopedTypeVariables++test-suite hbayesian-test+ import: warnings+ type: exitcode-stdio-1.0+ main-is: Main.hs+ other-modules:+ Test.Core+ Test.HHLO.Ops+ Test.HHLO.RNG+ Test.HHLO.Loops+ Test.MCMC+ Test.Examples+ Test.Chain+ Test.PPL+ Test.CorrelatedGaussian+ Common+ LinearRegressionRandomWalk+ GaussianProcessEllipticalSlice+ LogisticRegressionHMC+ BivariateGaussianMALA+ CorrelatedGaussianHMC+ build-depends:+ base >= 4.18.2.1 && < 4.20+ , hbayesian+ , hhlo >= 0.2.0.0 && < 0.3+ , text >= 2.0 && < 2.2+ , vector >= 0.13 && < 0.14+ , tasty >= 1.4 && < 1.6+ , tasty-hunit >= 0.10 && < 0.11+ , tasty-golden >= 2.3 && < 2.4+ , bytestring >= 0.11 && < 0.13+ hs-source-dirs: test, examples+ default-language: Haskell2010+ default-extensions:+ DataKinds+ KindSignatures+ TypeApplications+ TypeFamilies+ TypeOperators+ ScopedTypeVariables
+ src/HBayesian/Chain.hs view
@@ -0,0 +1,355 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Chain combinators for HBayesian v0.2.+--+-- This module hides the mechanical work of PJRT compilation, buffer+-- management, and host loops behind a simple configuration API.+--+-- Usage:+--+-- > ck <- compileSimpleKernel kernel logpdf+-- > (samples, diags) <- sampleChain ck [0.0, 0.0] $ burnIn 100 $ thin 2 $ defaultChainConfig { ccNumIterations = 1000 }+module HBayesian.Chain+ ( -- * Compiled kernels+ CompiledKernel+ , compileSimpleKernel+ , compileHMC+ -- * Chain configuration+ , ChainConfig (..)+ , defaultChainConfig+ , burnIn+ , thin+ , withSeed+ , verbose+ -- * Running chains+ , sampleChain+ , parallelChains+ , Diagnostic (..)+ ) where++import Control.Monad (when, zipWithM)+import Data.Proxy (Proxy (..))+import Data.Word (Word64)+import qualified Data.Vector.Storable as V++import HHLO.Core.Types+import HHLO.IR.AST (FuncArg (..), Module, TensorType)+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HHLO.Runtime.Buffer (toDevice, toDeviceF32, fromDeviceF32)+import HHLO.Runtime.Compile (compileWithOptions, defaultCompileOptions)+import HHLO.Runtime.Execute (execute)+import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTClient, PJRTExecutable, PJRTBuffer,+ bufferTypeU64)++import HBayesian.Core+import HBayesian.HHLO.Ops hiding (map)+import HBayesian.HHLO.PJRT+import HBayesian.MCMC.HMC (HMCState (..))++-----------------------------------------------------------------------------+-- CompiledKernel+-----------------------------------------------------------------------------++-- | Tag indicating which kind of step module was compiled.+data StepType = SimpleStep | HMCStep+ deriving (Eq, Show)++-- | A kernel that has been lowered to StableHLO modules but not yet+-- compiled to PJRT executables. The actual compilation happens inside+-- 'sampleChain' where the PJRT context is alive.+--+-- This design avoids the lifetime issue of PJRT handles: 'PJRTApi' and+-- 'PJRTClient' are raw pointers that become invalid when the plugin+-- is unloaded, so we cannot store compiled executables across calls.+data CompiledKernel = CompiledKernel+ { ckLdModule :: !Module+ , ckGradModule :: !(Maybe Module)+ , ckStepModule :: !Module+ , ckShape :: ![Int]+ , ckStepType :: !StepType+ }++-- | Render a module and compile it via PJRT.+compileModule :: PJRTApi -> PJRTClient -> Module -> IO PJRTExecutable+compileModule api client modl =+ compileWithOptions api client (render modl) defaultCompileOptions++-- | Create a shape list from a 'KnownShape' proxy.+shapeList :: forall s. KnownShape s => [Int]+shapeList = Prelude.map fromIntegral (shapeVal (Proxy @s))++-- | Create a 'TensorType' from shape/dtype proxies.+tensorTypeOf :: forall s d. (KnownShape s, KnownDType d) => TensorType+tensorTypeOf = tensorType (Proxy @s) (Proxy @d)++-- | UI64 key type.+keyType :: TensorType+keyType = tensorType (Proxy @'[2]) (Proxy @'UI64)++-----------------------------------------------------------------------------+-- Buffer helpers+-----------------------------------------------------------------------------++bufferFromF32 :: PJRTApi -> PJRTClient -> [Int] -> [Float] -> IO PJRTBuffer+bufferFromF32 api client dims vals =+ toDeviceF32 api client (V.fromList vals) (Prelude.map fromIntegral dims)++bufferFromUI64 :: PJRTApi -> PJRTClient -> [Int] -> [Word64] -> IO PJRTBuffer+bufferFromUI64 api client dims vals =+ toDevice api client (V.fromList vals) (Prelude.map fromIntegral dims) bufferTypeU64++bufferToF32 :: PJRTApi -> PJRTBuffer -> Int -> IO [Float]+bufferToF32 api buf n = V.toList <$> fromDeviceF32 api buf n++-----------------------------------------------------------------------------+-- Compiling SimpleKernel (RandomWalk, EllipticalSlice)+-----------------------------------------------------------------------------++-- | Build a 'CompiledKernel' from a 'SimpleKernel'.+--+-- The log-posterior is compiled to a separate module so the host+-- can recompute log-density between steps.+compileSimpleKernel :: forall s d.+ (KnownShape s, KnownDType d)+ => SimpleKernel s d+ -> (Tensor s d -> Builder (Tensor '[] d))+ -> CompiledKernel+compileSimpleKernel kernel logpdf =+ let ldMod = moduleFromBuilder @'[] @d "main"+ [FuncArg "theta" (tensorTypeOf @s @d)] $ do+ theta <- arg @s @d+ logpdf theta++ stepMod = moduleFromBuilder @s @d "main"+ [ FuncArg "key" keyType+ , FuncArg "pos" (tensorTypeOf @s @d)+ , FuncArg "ld" (tensorTypeOf @'[] @d)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @s @d+ ld <- arg @'[] @d+ (state', _info) <- kernelStep kernel (Key key) (State pos ld)+ return (statePosition state')+ in CompiledKernel ldMod Nothing stepMod (shapeList @s) SimpleStep++-----------------------------------------------------------------------------+-- Compiling HMC kernels (HMC, MALA)+-----------------------------------------------------------------------------++-- | Build a 'CompiledKernel' from an HMC-style kernel.+--+-- Requires both the log-posterior and its gradient.+compileHMC :: forall s d info.+ (KnownShape s, KnownDType d)+ => Kernel s d (HMCState s d) info+ -> (Tensor s d -> Builder (Tensor '[] d))+ -> Gradient s d+ -> CompiledKernel+compileHMC kernel logpdf grad =+ let ldMod = moduleFromBuilder @'[] @d "main"+ [FuncArg "theta" (tensorTypeOf @s @d)] $ do+ theta <- arg @s @d+ logpdf theta++ gradMod = moduleFromBuilder @s @d "main"+ [FuncArg "theta" (tensorTypeOf @s @d)] $ do+ theta <- arg @s @d+ grad theta++ stepMod = moduleFromBuilder @s @d "main"+ [ FuncArg "key" keyType+ , FuncArg "pos" (tensorTypeOf @s @d)+ , FuncArg "p" (tensorTypeOf @s @d)+ , FuncArg "ld" (tensorTypeOf @'[] @d)+ , FuncArg "g" (tensorTypeOf @s @d)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @s @d+ p <- arg @s @d+ ld <- arg @'[] @d+ g <- arg @s @d+ (state', _info) <- kernelStep kernel (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ in CompiledKernel ldMod (Just gradMod) stepMod (shapeList @s) HMCStep++-----------------------------------------------------------------------------+-- Chain configuration+-----------------------------------------------------------------------------++-- | Control parameters for a chain.+data ChainConfig = ChainConfig+ { ccNumIterations :: !Int+ , ccBurnIn :: !Int+ , ccThinning :: !Int+ , ccSeed :: !Word64+ , ccVerbose :: !Bool+ }++defaultChainConfig :: ChainConfig+defaultChainConfig = ChainConfig+ { ccNumIterations = 1000+ , ccBurnIn = 0+ , ccThinning = 1+ , ccSeed = 42+ , ccVerbose = False+ }++-- | Increase burn-in by N samples.+burnIn :: Int -> ChainConfig -> ChainConfig+burnIn n cfg = cfg { ccBurnIn = ccBurnIn cfg + n }++-- | Set thinning interval.+thin :: Int -> ChainConfig -> ChainConfig+thin n cfg = cfg { ccThinning = max 1 n }++-- | Override the PRNG seed.+withSeed :: Word64 -> ChainConfig -> ChainConfig+withSeed s cfg = cfg { ccSeed = s }++-- | Enable verbose progress output.+verbose :: ChainConfig -> ChainConfig+verbose cfg = cfg { ccVerbose = True }++-----------------------------------------------------------------------------+-- Diagnostics+-----------------------------------------------------------------------------++-- | A single-step diagnostic record.+data Diagnostic = Diagnostic+ { dStep :: !Int+ , dAccepted :: !Bool+ , dAcceptProb :: !Float+ }+ deriving (Show)++-----------------------------------------------------------------------------+-- Running a chain+-----------------------------------------------------------------------------++-- | Run a compiled kernel and return samples plus diagnostics.+--+-- The chain runs for @burnIn + numIterations * thinning@ steps total.+-- Samples are collected only after burn-in and only every @thinning@ steps.+--+-- This function opens a fresh PJRT context, compiles the modules,+-- executes the chain, and closes the context on return.+sampleChain :: CompiledKernel -> [Float] -> ChainConfig -> IO ([[Float]], [Diagnostic])+sampleChain ck pos0 cfg =+ withPJRTCPU $ \api client -> do+ let shape = ckShape ck+ nDim = product shape+ nTotal = ccNumIterations cfg+ nBurn = ccBurnIn cfg+ thinBy = ccThinning cfg+ seed = ccSeed cfg+ verb = ccVerbose cfg+ totalSteps = nBurn + nTotal * thinBy++ -- Compile modules inside the PJRT context+ ldExe <- compileModule api client (ckLdModule ck)+ gradExe <- case ckGradModule ck of+ Nothing -> return Nothing+ Just gm -> Just <$> compileModule api client gm+ stepExe <- compileModule api client (ckStepModule ck)++ -- Evaluate initial log-density+ posBuf0 <- bufferFromF32 api client shape pos0+ [ldBuf0] <- execute api ldExe [posBuf0]+ ld0 <- head <$> bufferToF32 api ldBuf0 1++ -- Evaluate initial gradient (if HMC)+ g0 <- case gradExe of+ Nothing -> return (replicate nDim 0.0)+ Just gE -> do+ [gBuf0] <- execute api gE [posBuf0]+ bufferToF32 api gBuf0 nDim++ -- Run the chain+ (positions, diags) <- runLoop api client stepExe ldExe gradExe+ (ckStepType ck) shape seed 0 pos0 ld0 g0 totalSteps verb++ -- Apply burn-in and thinning+ let postBurn = drop nBurn positions+ thinned = take nTotal $ every thinBy postBurn+ diagsPost = drop nBurn diags+ diagsThin = take nTotal $ every thinBy diagsPost++ return (thinned, diagsThin)+ where+ every n xs = case xs of+ [] -> []+ (y:ys) -> y : every n (drop (n - 1) ys)++-- | Run N independent chains in parallel.+--+-- Each chain gets a distinct PRNG seed and optionally a perturbed+-- initial position. Results are returned in the same order as seeds.+parallelChains :: Int -- ^ number of chains+ -> ([Float] -> [Float]) -- ^ perturbation for initial values+ -> CompiledKernel -- ^ compiled kernel+ -> [Float] -- ^ base initial position+ -> ChainConfig -- ^ chain configuration+ -> IO [([[Float]], [Diagnostic])]+parallelChains n perturb ck pos0 cfg =+ let seeds = [ccSeed cfg .. ccSeed cfg + fromIntegral n - 1]+ pos0s = pos0 : [perturb pos0 | _ <- [2..n]]+ in zipWithM (\s p -> sampleChain ck p (withSeed s cfg)) seeds pos0s++-- | Inner loop: run N steps, collecting all positions and diagnostics.+runLoop :: PJRTApi -> PJRTClient -> PJRTExecutable -> PJRTExecutable+ -> Maybe PJRTExecutable -> StepType -> [Int] -> Word64 -> Int+ -> [Float] -> Float -> [Float] -> Int -> Bool+ -> IO ([[Float]], [Diagnostic])+runLoop api client stepExe ldExe gradExe stepType shape seed step pos ld g n verb =+ go step pos ld g n []+ where+ nDim = product shape++ go _ _ _ _ 0 acc = return (reverse (map fst acc), reverse (map snd acc))+ go st p l gr remaining acc = do+ let key = [seed, fromIntegral st]++ keyBuf <- bufferFromUI64 api client [2] key+ posBuf <- bufferFromF32 api client shape p++ (newPos, newLD, newG, acceptProb) <- case stepType of+ SimpleStep -> do+ ldBuf <- bufferFromF32 api client [] [l]+ [newPosBuf] <- execute api stepExe [keyBuf, posBuf, ldBuf]+ newPos <- bufferToF32 api newPosBuf nDim+ -- Recompute log-density for next iteration+ [newLdBuf] <- execute api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ let changed = newPos /= p+ let prob = if changed then 1.0 else 0.0+ return (newPos, newLd, replicate nDim 0.0, prob)++ HMCStep -> do+ pBuf <- bufferFromF32 api client shape (replicate nDim 0.0)+ ldBuf <- bufferFromF32 api client [] [l]+ gBuf <- bufferFromF32 api client shape gr+ [newPosBuf] <- execute api stepExe [keyBuf, posBuf, pBuf, ldBuf, gBuf]+ newPos <- bufferToF32 api newPosBuf nDim+ -- Recompute log-density and gradient for next iteration+ [newLdBuf] <- execute api ldExe [newPosBuf]+ [newLd] <- bufferToF32 api newLdBuf 1+ newG <- case gradExe of+ Just gE -> do+ [newGBuf] <- execute api gE [newPosBuf]+ bufferToF32 api newGBuf nDim+ Nothing -> return (replicate nDim 0.0)+ let changed = newPos /= p+ let prob = if changed then 1.0 else 0.0+ return (newPos, newLd, newG, prob)++ when verb $ do+ putStrLn $ "Step " ++ show st ++ ": pos=" ++ show newPos+ ++ " accept=" ++ show acceptProb++ let diag = Diagnostic st (acceptProb > 0.5) acceptProb+ go (st + 1) newPos newLD newG (remaining - 1) ((newPos, diag) : acc)
+ src/HBayesian/Core.hs view
@@ -0,0 +1,74 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}++-- | Core types and abstractions for HBayesian.+module HBayesian.Core+ ( -- * PRNG+ Key (..)+ -- * Gradients+ , Gradient+ -- * State and diagnostics+ , State (..)+ , Info (..)+ -- * Kernel abstraction+ , Kernel (..)+ , SimpleKernel+ ) where++import HHLO.Core.Types (DType (..), Shape)+import HHLO.IR.Builder (Builder, Tensor)++--------------------------------------------------------------------------------+-- PRNG+--------------------------------------------------------------------------------++-- | A functional PRNG key backed by Threefry.+newtype Key = Key { unKey :: Tensor '[2] 'UI64 }++--------------------------------------------------------------------------------+-- Gradients+--------------------------------------------------------------------------------++-- | A user-provided gradient function.+type Gradient (s :: Shape) (d :: DType) = Tensor s d -> Builder (Tensor s d)++--------------------------------------------------------------------------------+-- State+--------------------------------------------------------------------------------++-- | Minimal algorithm state: position and log-density.+data State (s :: Shape) (d :: DType) = State+ { statePosition :: !(Tensor s d)+ , stateLogDensity :: !(Tensor '[] d)+ }++--------------------------------------------------------------------------------+-- Diagnostics+--------------------------------------------------------------------------------++-- | Diagnostic information produced by a single transition.+data Info (s :: Shape) (d :: DType) = Info+ { infoAcceptProb :: !(Tensor '[] d)+ , infoAccepted :: !(Tensor '[] 'Bool)+ , infoNumSteps :: !(Tensor '[] 'I64)+ }++--------------------------------------------------------------------------------+-- Kernel+--------------------------------------------------------------------------------++-- | A transition kernel, polymorphic in state and info.+--+-- Algorithms define their own state and info types, then expose a+-- 'Kernel s d state info'. The convenience alias 'SimpleKernel' covers+-- the common gradient-free case.+data Kernel (s :: Shape) (d :: DType) state info = Kernel+ { kernelInit :: !(Key -> Tensor s d -> Builder state)+ , kernelStep :: !(Key -> state -> Builder (state, info))+ }++-- | A kernel that uses the base 'State' and 'Info' types.+type SimpleKernel s d = Kernel s d (State s d) (Info s d)
+ src/HBayesian/Diagnostics.hs view
@@ -0,0 +1,113 @@+{-# LANGUAGE ScopedTypeVariables #-}++-- | Diagnostics for MCMC chains.+--+-- These functions operate on host-side 'Diagnostic' records collected+-- by 'HBayesian.Chain.sampleChain'.+module HBayesian.Diagnostics+ ( -- * Acceptance diagnostics+ acceptanceRate+ , meanAcceptProb+ -- * Convergence diagnostics+ , rHat+ -- * Effective sample size+ , ess+ ) where++import HBayesian.Chain (Diagnostic (..))++-----------------------------------------------------------------------------+-- Acceptance diagnostics+-----------------------------------------------------------------------------++-- | Fraction of steps that were accepted.+acceptanceRate :: [Diagnostic] -> Double+acceptanceRate diags =+ let accepted = length (filter dAccepted diags)+ total = length diags+ in if total == 0 then 0.0 else fromIntegral accepted / fromIntegral total++-- | Mean acceptance probability across all steps.+meanAcceptProb :: [Diagnostic] -> Double+meanAcceptProb diags =+ let total = length diags+ in if total == 0 then 0.0+ else sum (map (realToFrac . dAcceptProb) diags) / fromIntegral total++-----------------------------------------------------------------------------+-- Convergence diagnostics+-----------------------------------------------------------------------------++-- | Potential scale reduction factor (R-hat).+--+-- Requires at least two chains. Returns 'Nothing' if insufficient data.+rHat :: [[Diagnostic]] -> Maybe Double+rHat chains+ | length chains < 2 = Nothing+ | any null chains = Nothing+ | otherwise = Just $ computeRHat (map (map dAcceptProb) chains)++computeRHat :: [[Float]] -> Double+computeRHat chains =+ let n = fromIntegral (length (head chains)) :: Double+ -- Within-chain variance+ chainVars = map varianceF chains+ w = meanD chainVars+ -- Between-chain variance+ b = n * varianceD (map meanF chains)+ -- Pooled variance+ vHat = ((n - 1) / n) * w + (1 / n) * b+ in sqrt (vHat / w)++meanF :: [Float] -> Double+meanF xs = sum (map realToFrac xs) / fromIntegral (length xs)++meanD :: [Double] -> Double+meanD xs = sum xs / fromIntegral (length xs)++varianceF :: [Float] -> Double+varianceF xs =+ let m = meanF xs+ n = fromIntegral (length xs) :: Double+ in if n <= 1 then 0.0+ else sum (map (\x -> (realToFrac x - m) ** 2) xs) / (n - 1)++varianceD :: [Double] -> Double+varianceD xs =+ let m = meanD xs+ n = fromIntegral (length xs) :: Double+ in if n <= 1 then 0.0+ else sum (map (\x -> (x - m) ** 2) xs) / (n - 1)++-----------------------------------------------------------------------------+-- Effective sample size+-----------------------------------------------------------------------------++-- | Effective sample size estimate (naive autocorrelation method).+--+-- This is a simplified ESS that uses the autocorrelation of the+-- acceptance probabilities as a proxy for chain mixing.+ess :: [Diagnostic] -> Double+ess diags+ | null diags = 0.0+ | otherwise =+ let n = fromIntegral (length diags)+ vals = map (realToFrac . dAcceptProb) diags+ rhoK = takeWhile (> 0.05) (autocorrelations vals)+ in n / (1 + 2 * sum rhoK)++-- | Autocorrelation at lag k.+autocorrelations :: [Double] -> [Double]+autocorrelations xs =+ let n = length xs+ m = meanD xs+ v = varianceD xs+ centred = map (\x -> x - m) xs+ in if v == 0 then []+ else [ lagKAutocorr centred k / v | k <- [1 .. n `div` 2] ]++lagKAutocorr :: [Double] -> Int -> Double+lagKAutocorr xs k =+ let n = length xs+ in if k >= n then 0.0+ else sum [ xs !! i * xs !! (i + k) | i <- [0 .. n - k - 1] ] / fromIntegral (n - k)
+ src/HBayesian/HHLO/Compile.hs view
@@ -0,0 +1,32 @@+{-# LANGUAGE OverloadedStrings #-}++-- | Compilation and execution utilities for HBayesian kernels.+--+-- This module provides thin wrappers around HHLO's PJRT runtime.+-- In Phase 1 it is intentionally minimal; it will expand as the project+-- matures.+module HBayesian.HHLO.Compile+ ( compileModule+ , renderBuilder+ ) where++import qualified Data.Text as T+import HHLO.IR.AST (Module)+import HHLO.IR.Builder (Builder)+import HHLO.IR.Pretty (render)+import qualified HHLO.Runtime.Compile as RT+import qualified HHLO.Runtime.PJRT.Plugin as RT+import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTClient, PJRTExecutable)++-- | Render a 'Builder' action to its StableHLO MLIR text.+--+-- This is primarily useful for debugging and golden tests.+renderBuilder :: Builder a -> Module -> T.Text+renderBuilder _ = render++-- | Compile a 'Module' to a PJRT executable.+--+-- Requires an active PJRT API and client. Use 'RT.withPJRT' to obtain them.+compileModule :: PJRTApi -> PJRTClient -> Module -> IO PJRTExecutable+compileModule api client modl =+ RT.compileWithOptions api client (render modl) RT.defaultCompileOptions
+ src/HBayesian/HHLO/Loops.hs view
@@ -0,0 +1,182 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Multi-value control flow primitives for Bayesian state patterns.+--+-- HHLO 0.2.0.0 provides 'whileLoop2' and 'whileLoopN'.+-- This module adds 'whileLoop3', 'whileLoop4', and 'conditional3', 'conditional4'+-- as mechanical boilerplate.+module HBayesian.HHLO.Loops+ ( whileLoop3+ , whileLoop4+ , conditional3+ , conditional4+ ) where++import Data.Proxy+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (Region(..))+import HHLO.IR.Builder++-----------------------------------------------------------------------------+-- whileLoop3+-----------------------------------------------------------------------------++whileLoop3 :: forall s1 d1 s2 d2 s3 d3.+ ( KnownShape s1, KnownDType d1+ , KnownShape s2, KnownDType d2+ , KnownShape s3, KnownDType d3+ )+ => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3+ -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor '[] 'Bool))+ -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3))+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3)+whileLoop3 init1 init2 init3 cond body = do+ let initVids = [tensorValue init1, tensorValue init2, tensorValue init3]+ initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+ , tensorType (Proxy @s2) (Proxy @d2)+ , tensorType (Proxy @s3) (Proxy @d3)+ ]+ boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++ -- Build cond region+ condBlock <- runBlockBuilder initTypes $ do+ v1 <- arg @s1 @d1+ v2 <- arg @s2 @d2+ v3 <- arg @s3 @d3+ c <- cond v1 v2 v3+ emitReturn [tensorValue c] [boolType]++ -- Build body region+ bodyBlock <- runBlockBuilder initTypes $ do+ v1 <- arg @s1 @d1+ v2 <- arg @s2 @d2+ v3 <- arg @s3 @d3+ (r1, r2, r3) <- body v1 v2 v3+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] initTypes++ vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+ [Region [condBlock], Region [bodyBlock]] initTypes+ case vids of+ [vid1, vid2, vid3] -> return (Tensor vid1, Tensor vid2, Tensor vid3)+ _ -> error "whileLoop3: expected exactly three results"++-----------------------------------------------------------------------------+-- whileLoop4+-----------------------------------------------------------------------------++whileLoop4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+ ( KnownShape s1, KnownDType d1+ , KnownShape s2, KnownDType d2+ , KnownShape s3, KnownDType d3+ , KnownShape s4, KnownDType d4+ )+ => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4+ -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tensor '[] 'Bool))+ -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3, Tensor s4 d4))+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3, Tensor s4 d4)+whileLoop4 init1 init2 init3 init4 cond body = do+ let initVids = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4]+ initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+ , tensorType (Proxy @s2) (Proxy @d2)+ , tensorType (Proxy @s3) (Proxy @d3)+ , tensorType (Proxy @s4) (Proxy @d4)+ ]+ boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++ condBlock <- runBlockBuilder initTypes $ do+ v1 <- arg @s1 @d1+ v2 <- arg @s2 @d2+ v3 <- arg @s3 @d3+ v4 <- arg @s4 @d4+ c <- cond v1 v2 v3 v4+ emitReturn [tensorValue c] [boolType]++ bodyBlock <- runBlockBuilder initTypes $ do+ v1 <- arg @s1 @d1+ v2 <- arg @s2 @d2+ v3 <- arg @s3 @d3+ v4 <- arg @s4 @d4+ (r1, r2, r3, r4) <- body v1 v2 v3 v4+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] initTypes++ vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+ [Region [condBlock], Region [bodyBlock]] initTypes+ case vids of+ [vid1, vid2, vid3, vid4] -> return (Tensor vid1, Tensor vid2, Tensor vid3, Tensor vid4)+ _ -> error "whileLoop4: expected exactly four results"++-----------------------------------------------------------------------------+-- conditional3+-----------------------------------------------------------------------------++conditional3 :: forall s1 d1 s2 d2 s3 d3.+ ( KnownShape s1, KnownDType d1+ , KnownShape s2, KnownDType d2+ , KnownShape s3, KnownDType d3+ )+ => Tensor '[] 'Bool+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3)+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3)+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3)+conditional3 p trueThunk falseThunk = do+ let types = [ tensorType (Proxy @s1) (Proxy @d1)+ , tensorType (Proxy @s2) (Proxy @d2)+ , tensorType (Proxy @s3) (Proxy @d3)+ ]+ boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++ trueBlock <- runBlockBuilder [] $ do+ (r1, r2, r3) <- trueThunk+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types++ falseBlock <- runBlockBuilder [] $ do+ (r1, r2, r3) <- falseThunk+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types++ let (Tensor predVid) = p+ vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+ [Region [trueBlock], Region [falseBlock]] types+ case vids of+ [vid1, vid2, vid3] -> return (Tensor vid1, Tensor vid2, Tensor vid3)+ _ -> error "conditional3: expected exactly three results"++-----------------------------------------------------------------------------+-- conditional4+-----------------------------------------------------------------------------++conditional4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+ ( KnownShape s1, KnownDType d1+ , KnownShape s2, KnownDType d2+ , KnownShape s3, KnownDType d3+ , KnownShape s4, KnownDType d4+ )+ => Tensor '[] 'Bool+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3, Tensor s4 d4)+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3, Tensor s4 d4)+ -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3, Tensor s4 d4)+conditional4 p trueThunk falseThunk = do+ let types = [ tensorType (Proxy @s1) (Proxy @d1)+ , tensorType (Proxy @s2) (Proxy @d2)+ , tensorType (Proxy @s3) (Proxy @d3)+ , tensorType (Proxy @s4) (Proxy @d4)+ ]+ boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++ trueBlock <- runBlockBuilder [] $ do+ (r1, r2, r3, r4) <- trueThunk+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types++ falseBlock <- runBlockBuilder [] $ do+ (r1, r2, r3, r4) <- falseThunk+ emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types++ let (Tensor predVid) = p+ vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+ [Region [trueBlock], Region [falseBlock]] types+ case vids of+ [vid1, vid2, vid3, vid4] -> return (Tensor vid1, Tensor vid2, Tensor vid3, Tensor vid4)+ _ -> error "conditional4: expected exactly four results"
+ src/HBayesian/HHLO/Ops.hs view
@@ -0,0 +1,289 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.HHLO.Ops+ ( module HHLO.EDSL.Ops+ , module HHLO.Core.Types+ -- * Missing primitive ops+ , sqrt'+ , rsqrt'+ , sin'+ , cos'+ , tan'+ , pow'+ , log1p'+ , floor'+ , ceil'+ -- * Element-wise comparison (returns tensor of bools)+ , lessThanEW+ , greaterThanEW+ , equalEW+ -- * Convenience aliases+ , tadd+ , tsub+ , tmul+ , tdiv+ , tnegate+ , tabs+ , texp+ , tlog+ , tsqrt+ , trsqrt+ , tsin+ , tcos+ , ttan+ , tpow+ , tlessThan+ , tminimum+ , tmaximum+ , tselect+ , tconstant+ , tsumAll+ , treshape+ , tbroadcast+ , tslice1+ , tpack2+ , tpack3+ , tsigmoid+ ) where++import Data.Int (Int64)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (Attribute (..))+import HHLO.IR.Builder+import Prelude hiding (negate, minimum, maximum)++-----------------------------------------------------------------------------+-- Missing primitive ops (not in HHLO 0.2.0.0)+-----------------------------------------------------------------------------++sqrt' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+sqrt' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.sqrt" [x] [ttype] [] ttype+ return (Tensor vid)++rsqrt' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+rsqrt' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.rsqrt" [x] [ttype] [] ttype+ return (Tensor vid)++sin' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+sin' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.sine" [x] [ttype] [] ttype+ return (Tensor vid)++cos' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+cos' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.cosine" [x] [ttype] [] ttype+ return (Tensor vid)++tan' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tan' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.tangent" [x] [ttype] [] ttype+ return (Tensor vid)++pow' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s d)+pow' (Tensor x) (Tensor y) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.power" [x, y] [ttype, ttype] [] ttype+ return (Tensor vid)++log1p' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+log1p' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.log_plus_one" [x] [ttype] [] ttype+ return (Tensor vid)++floor' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+floor' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.floor" [x] [ttype] [] ttype+ return (Tensor vid)++ceil' :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+ceil' (Tensor x) = do+ let ttype = tensorType (Proxy @s) (Proxy @d)+ vid <- emitOp "stablehlo.ceil" [x] [ttype] [] ttype+ return (Tensor vid)++-----------------------------------------------------------------------------+-- Element-wise comparison (HHLO's compare returns scalar; we fix it)+-----------------------------------------------------------------------------++lessThanEW :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+lessThanEW (Tensor x) (Tensor y) = do+ let inType = tensorType (Proxy @s) (Proxy @d)+ outType = tensorType (Proxy @s) (Proxy @'Bool)+ vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]+ [AttrString "comparison_direction" "LT"] outType+ return (Tensor vid)++greaterThanEW :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+greaterThanEW (Tensor x) (Tensor y) = do+ let inType = tensorType (Proxy @s) (Proxy @d)+ outType = tensorType (Proxy @s) (Proxy @'Bool)+ vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]+ [AttrString "comparison_direction" "GT"] outType+ return (Tensor vid)++equalEW :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+equalEW (Tensor x) (Tensor y) = do+ let inType = tensorType (Proxy @s) (Proxy @d)+ outType = tensorType (Proxy @s) (Proxy @'Bool)+ vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]+ [AttrString "comparison_direction" "EQ"] outType+ return (Tensor vid)++-----------------------------------------------------------------------------+-- Convenience aliases+-----------------------------------------------------------------------------++tadd, tsub, tmul, tdiv :: forall s d.+ (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s d)+tadd = add+tsub = sub+tmul = multiply+tdiv = divide++tnegate :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tnegate = negate++tabs :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tabs = abs'++texp :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+texp = exponential++tlog :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tlog = logarithm++tsqrt :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tsqrt = sqrt'++trsqrt :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+trsqrt = rsqrt'++tsin :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tsin = sin'++tcos :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+tcos = cos'++ttan :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor s d)+ttan = tan'++tpow :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s d)+tpow = pow'++tlessThan :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor '[] 'Bool)+tlessThan = lessThan++tselect :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s 'Bool -> Tensor s d -> Tensor s d -> Builder (Tensor s d)+tselect = select++tconstant :: forall s d. (KnownShape s, KnownDType d)+ => Double -> Builder (Tensor s d)+tconstant = constant++tsumAll :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Builder (Tensor '[] d)+tsumAll = reduceSum++treshape :: forall sFrom sTo d.+ (KnownShape sFrom, KnownShape sTo, KnownDType d)+ => Tensor sFrom d -> Builder (Tensor sTo d)+treshape = reshape++tbroadcast :: forall sFrom sTo d.+ (KnownShape sFrom, KnownShape sTo, KnownDType d)+ => [Int64] -> Tensor sFrom d -> Builder (Tensor sTo d)+tbroadcast = broadcastWithDims++tminimum :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s d)+tminimum = minimum++tmaximum :: forall s d. (KnownShape s, KnownDType d)+ => Tensor s d -> Tensor s d -> Builder (Tensor s d)+tmaximum = maximum++-----------------------------------------------------------------------------+-- Slice / pack helpers for parameter vectors+-----------------------------------------------------------------------------++-- | Extract a single scalar element from a 1-D tensor at a constant index.+tslice1 :: forall n d. (KnownShape '[n], KnownDType d)+ => Tensor '[n] d -> Int64 -> Builder (Tensor '[] d)+tslice1 vec i = do+ sliced <- slice @'[n] @'[1] @d vec [i] [i+1] [1]+ treshape @'[1] @'[] sliced++-- | Pack two scalar tensors into a rank-1 tensor of shape [2].+tpack2 :: forall d. (KnownDType d)+ => Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[2] d)+tpack2 x y = do+ x1 <- treshape @'[] @'[1] x+ y1 <- treshape @'[] @'[1] y+ concatenate @'[1] @'[2] @d 0 [x1, y1]++-- | Pack three scalar tensors into a rank-1 tensor of shape [3].+tpack3 :: forall d. (KnownDType d)+ => Tensor '[] d -> Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[3] d)+tpack3 x y z = do+ x1 <- treshape @'[] @'[1] x+ y1 <- treshape @'[] @'[1] y+ z1 <- treshape @'[] @'[1] z+ concatenate @'[1] @'[3] @d 0 [x1, y1, z1]++-----------------------------------------------------------------------------+-- Sigmoid and matmul+-----------------------------------------------------------------------------++-- | Element-wise sigmoid: 1 / (1 + exp(-x)).+tsigmoid :: forall s. (KnownShape s)+ => Tensor s 'F32 -> Builder (Tensor s 'F32)+tsigmoid x = do+ negX <- tnegate x+ expNegX <- texp negX+ one <- tconstant @s @'F32 1.0+ denom <- tadd one expNegX+ tdiv one denom++
+ src/HBayesian/HHLO/PJRT.hs view
@@ -0,0 +1,60 @@+{-# LANGUAGE OverloadedStrings #-}++-- | PJRT plugin discovery and execution helpers.+--+-- This module mirrors HHLO's strategy: a 'pjrt_script.sh' downloads the+-- plugin into 'deps/pjrt/', and the Haskell code references that path by+-- default. Users can override with the 'HBAYESIAN_PJRT_PLUGIN' environment+-- variable.+module HBayesian.HHLO.PJRT+ ( getPluginPath+ , withPJRTCPU+ , withPJRTCPUAt+ ) where++import System.Directory (doesFileExist)+import System.Environment (lookupEnv)+import HHLO.Runtime.PJRT.Plugin (withPJRT)+import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTClient)++-- | Return the path to the PJRT CPU plugin.+--+-- Priority:+-- 1. @HBAYESIAN_PJRT_PLUGIN@ environment variable+-- 2. @deps/pjrt/libpjrt_cpu.so@ (downloaded by 'scripts/pjrt_script.sh')+-- 3. Runtime error with instructions+getPluginPath :: IO FilePath+getPluginPath = do+ mEnv <- lookupEnv "HBAYESIAN_PJRT_PLUGIN"+ case mEnv of+ Just p -> return p+ Nothing -> do+ let defaultPath = "deps/pjrt/libpjrt_cpu.so"+ exists <- doesFileExist defaultPath+ if exists+ then return defaultPath+ else error $ unlines+ [ "PJRT CPU plugin not found at: " ++ defaultPath+ , ""+ , "To fix this, either:"+ , " 1. Run the download script:"+ , " ./scripts/pjrt_script.sh"+ , " 2. Set the environment variable to an existing plugin:"+ , " export HBAYESIAN_PJRT_PLUGIN=/path/to/libpjrt_cpu.so"+ ]++-- | Bracket-style PJRT initialization using a specific plugin path.+withPJRTCPUAt :: FilePath+ -> (HHLO.Runtime.PJRT.Types.PJRTApi -> HHLO.Runtime.PJRT.Types.PJRTClient -> IO a)+ -> IO a+withPJRTCPUAt = withPJRT++-- | Bracket-style PJRT initialization using the CPU plugin.+--+-- Equivalent to 'HHLO.Runtime.PJRT.Plugin.withPJRTCPU' but respects the+-- 'HBAYESIAN_PJRT_PLUGIN' environment variable and our 'deps/pjrt/'+-- default path.+withPJRTCPU :: (HHLO.Runtime.PJRT.Types.PJRTApi -> HHLO.Runtime.PJRT.Types.PJRTClient -> IO a) -> IO a+withPJRTCPU action = do+ path <- getPluginPath+ withPJRT path action
+ src/HBayesian/HHLO/RNG.hs view
@@ -0,0 +1,96 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.HHLO.RNG+ ( -- Re-export Key from Core so all modules use the same type+ Key (..)+ , splitKey+ , rngUniformF32+ , rngUniformF64+ , rngNormalF32+ , rngNormalF64+ , rngBernoulli+ ) where++import Data.Word (Word64)+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HBayesian.Core (Key (..))+import HBayesian.HHLO.Ops++-----------------------------------------------------------------------------+-- Key splitting+-----------------------------------------------------------------------------++splitKey :: Key -> Builder (Key, Key)+splitKey (Key k) = do+ (k1, bits) <- rngBitGenerator @'[2] k+ return (Key k1, Key bits)++-----------------------------------------------------------------------------+-- Uniform [0,1)+-----------------------------------------------------------------------------++rngUniformF32 :: forall s. KnownShape s => Key -> Builder (Tensor s 'F32)+rngUniformF32 (Key k) = do+ (_, bits) <- rngBitGenerator @s k+ bitsF32 <- convert bits+ maxVal <- constant @'[] @'F32 (fromIntegral (maxBound :: Word64))+ maxValBC <- broadcastWithDims @'[] @s [] maxVal+ tdiv bitsF32 maxValBC++rngUniformF64 :: forall s. KnownShape s => Key -> Builder (Tensor s 'F64)+rngUniformF64 (Key k) = do+ (_, bits) <- rngBitGenerator @s k+ bitsF64 <- convert bits+ maxVal <- constant @'[] @'F64 (fromIntegral (maxBound :: Word64))+ maxValBC <- broadcastWithDims @'[] @s [] maxVal+ tdiv bitsF64 maxValBC++-----------------------------------------------------------------------------+-- Standard normal (Box-Muller)+-----------------------------------------------------------------------------++rngNormalF32 :: forall s. KnownShape s => Key -> Builder (Tensor s 'F32)+rngNormalF32 key = do+ (key1, key2) <- splitKey key+ u1 <- rngUniformF32 key1+ u2 <- rngUniformF32 key2+ twoPi <- constant @'[] @'F32 (2.0 * pi)+ negTwo <- constant @'[] @'F32 (-2.0)+ twoPiBC <- broadcastWithDims @'[] @s [] twoPi+ negTwoBC <- broadcastWithDims @'[] @s [] negTwo+ logU1 <- tlog u1+ term1 <- tmul negTwoBC logU1+ sqrtTerm1 <- tsqrt term1+ angle <- tmul twoPiBC u2+ cosAngle <- tcos angle+ tmul sqrtTerm1 cosAngle++rngNormalF64 :: forall s. KnownShape s => Key -> Builder (Tensor s 'F64)+rngNormalF64 key = do+ (key1, key2) <- splitKey key+ u1 <- rngUniformF64 key1+ u2 <- rngUniformF64 key2+ twoPi <- constant @'[] @'F64 (2.0 * pi)+ negTwo <- constant @'[] @'F64 (-2.0)+ twoPiBC <- broadcastWithDims @'[] @s [] twoPi+ negTwoBC <- broadcastWithDims @'[] @s [] negTwo+ logU1 <- tlog u1+ term1 <- tmul negTwoBC logU1+ sqrtTerm1 <- tsqrt term1+ angle <- tmul twoPiBC u2+ cosAngle <- tcos angle+ tmul sqrtTerm1 cosAngle++-----------------------------------------------------------------------------+-- Bernoulli+-----------------------------------------------------------------------------++rngBernoulli :: forall s. KnownShape s => Key -> Tensor s 'F32 -> Builder (Tensor s 'Bool)+rngBernoulli key probs = do+ u <- rngUniformF32 key+ lessThanEW u probs
+ src/HBayesian/InferenceLoop.hs view
@@ -0,0 +1,68 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | Host-side inference loops for running compiled kernels.+--+-- Phase 2 uses a hybrid strategy: the kernel step is compiled to a PJRT+-- executable and executed repeatedly from Haskell IO. This avoids the+-- complexity of compiling the entire chain into a single XLA graph while+-- still running the heavy math on the device.+module HBayesian.InferenceLoop+ ( InferenceConfig (..)+ , defaultInferenceConfig+ , sampleChain+ ) where++import Data.Int (Int64)+import Data.Vector.Storable (Vector)+import qualified Data.Vector.Storable as V+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import qualified HHLO.Runtime.Buffer as RTBuf+import qualified HHLO.Runtime.Compile as RT+import qualified HHLO.Runtime.Execute as RTExec+import qualified HHLO.Runtime.PJRT.Plugin as RT+import HBayesian.Core+import HBayesian.HHLO.Compile+import HBayesian.HHLO.RNG ()++-- | Configuration for the inference loop.+data InferenceConfig = InferenceConfig+ { icNumWarmup :: Int -- ^ number of warmup (burn-in) steps+ , icNumSamples :: Int -- ^ number of sampling steps+ , icThinning :: Int -- ^ keep every N-th sample (1 = no thinning)+ }++defaultInferenceConfig :: InferenceConfig+defaultInferenceConfig = InferenceConfig+ { icNumWarmup = 0+ , icNumSamples = 1000+ , icThinning = 1+ }++-- | Run a single chain and collect samples.+--+-- This is a stub for Phase 2. A full implementation requires PJRT+-- plugin installation and buffer management. The function demonstrates+-- the intended API and will be completed when PJRT is available.+sampleChain :: forall s d state info.+ (KnownShape s, KnownDType d)+ => InferenceConfig+ -> Kernel s d state info+ -> Key -- ^ initial PRNG key+ -> Vector Double -- ^ initial position (host)+ -> IO [Vector Double] -- ^ collected samples (host)+sampleChain config kernel key0 pos0 = do+ -- NOTE: Full PJRT-based execution requires the PJRT CPU plugin.+ -- For Phase 2, this function demonstrates the API shape.+ -- When PJRT is available, the implementation will:+ -- 1. Compile kernelInit and kernelStep to PJRT executables+ -- 2. Transfer key and position to device buffers+ -- 3. Run the warmup + sampling loop+ -- 4. Read back samples to host vectors+ putStrLn "sampleChain: PJRT-based execution not yet enabled in Phase 2"+ return [pos0]
+ src/HBayesian/MCMC/EllipticalSlice.hs view
@@ -0,0 +1,58 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.MCMC.EllipticalSlice+ ( ellipticalSlice+ ) where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import qualified HBayesian.HHLO.RNG as RNG++ellipticalSlice :: forall s d.+ (KnownShape s, KnownDType d)+ => (Tensor s d -> Builder (Tensor '[] d))+ -> SimpleKernel s d+ellipticalSlice logpdf = Kernel { kernelInit = kernelInit, kernelStep = kernelStep }+ where+ kernelInit _key pos = do+ ld <- logpdf pos+ return $ State pos ld++ kernelStep key state = do+ (key1, key2) <- RNG.splitKey key+ let pos = statePosition state+ let ld = stateLogDensity state++ nu <- RNG.rngNormalF32 key1 >>= convert @s @'F32 @d++ thetaRaw <- RNG.rngUniformF32 key2 >>= convert @'[] @'F32 @d+ twoPi <- constant @'[] @d (2.0 * pi)+ theta <- tmul thetaRaw twoPi++ c <- tcos theta+ sTheta <- tsin theta+ cBC <- tbroadcast @'[] @s [] c+ sBC <- tbroadcast @'[] @s [] sTheta+ pc <- tmul pos cBC+ ns <- tmul nu sBC+ pos' <- tadd pc ns++ ld' <- logpdf pos'++ zero <- constant @'[] @d 0.0+ diff <- tsub ld' ld+ accept <- tlessThan zero diff+ acceptS <- tbroadcast @'[] @s [] accept++ newPos <- tselect acceptS pos' pos+ newLd <- tselect accept ld' ld++ one <- constant @'[] @'I64 1+ infoAcceptProb <- texp diff+ let info = Info infoAcceptProb accept one+ return (State newPos newLd, info)
+ src/HBayesian/MCMC/HMC.hs view
@@ -0,0 +1,113 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.MCMC.HMC+ ( HMCConfig (..)+ , HMCState (..)+ , hmc+ ) where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import qualified HBayesian.HHLO.RNG as RNG++-- | Configuration for Hamiltonian Monte Carlo.+data HMCConfig = HMCConfig+ { hmcStepSize :: Double+ , hmcNumLeapfrogSteps :: Int+ }++-- | HMC-specific state.+data HMCState (s :: Shape) (d :: DType) = HMCState+ { hmcPosition :: !(Tensor s d)+ , hmcMomentum :: !(Tensor s d)+ , hmcLogDens :: !(Tensor '[] d)+ , hmcGradient :: !(Tensor s d)+ }++hmc :: forall s d.+ (KnownShape s, KnownDType d)+ => (Tensor s d -> Builder (Tensor '[] d))+ -> Gradient s d+ -> HMCConfig+ -> Kernel s d (HMCState s d) (Info s d)+hmc logpdf grad config = Kernel { kernelInit = kernelInit, kernelStep = kernelStep }+ where+ epsVal = hmcStepSize config+ nSteps = hmcNumLeapfrogSteps config++ kernelInit _key pos = do+ ld <- logpdf pos+ g <- grad pos+ zeroM <- tconstant 0.0+ return $ HMCState pos zeroM ld g++ kernelStep key state = do+ (key1, key2) <- RNG.splitKey key+ let pos0 = hmcPosition state+ let g0 = hmcGradient state+ let ld0 = hmcLogDens state++ p0 <- RNG.rngNormalF32 key1 >>= convert @s @'F32 @d++ currentK <- do+ pSq <- tmul p0 p0+ pSum <- tsumAll pSq+ half <- constant @'[] @d 0.5+ tmul half pSum+ currentH <- tsub currentK ld0++ (pos', p', g', ld') <- leapfrog pos0 p0 g0++ proposedK <- do+ pSq <- tmul p' p'+ pSum <- tsumAll pSq+ half <- constant @'[] @d 0.5+ tmul half pSum+ proposedH <- tsub proposedK ld'++ logAccept <- tsub currentH proposedH+ zero <- constant @'[] @d 0.0+ logAlpha <- tminimum logAccept zero++ u <- RNG.rngUniformF32 key2 >>= convert @'[] @'F32 @d+ logU <- tlog u+ accept <- tlessThan logU logAlpha+ acceptS <- tbroadcast @'[] @s [] accept++ newPos <- tselect acceptS pos' pos0+ newP <- tselect acceptS p' p0+ newG <- tselect acceptS g' g0+ newLd <- tselect accept ld' ld0++ infoAcceptProb <- texp logAlpha+ infoNumSteps <- constant @'[] @'I64 (fromIntegral nSteps)+ let info = Info infoAcceptProb accept infoNumSteps+ return (HMCState newPos newP newLd newG, info)++ leapfrog :: Tensor s d -> Tensor s d -> Tensor s d+ -> Builder (Tensor s d, Tensor s d, Tensor s d, Tensor '[] d)+ leapfrog pos0 p0 g0 = go nSteps pos0 p0 g0+ where+ go 0 pos p g = do+ ld <- logpdf pos+ return (pos, p, g, ld)+ go k pos p g = do+ halfEps <- constant @s @d (epsVal / 2.0)+ gScaled <- tmul g halfEps+ pHalf <- tadd p gScaled++ epsT <- constant @s @d epsVal+ pScaled <- tmul pHalf epsT+ pos' <- tadd pos pScaled++ g' <- grad pos'++ gScaled' <- tmul g' halfEps+ p' <- tadd pHalf gScaled'++ go (k-1) pos' p' g'
+ src/HBayesian/MCMC/MALA.hs view
@@ -0,0 +1,31 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.MCMC.MALA+ ( MALAConfig (..)+ , mala+ ) where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import HBayesian.HHLO.RNG+import HBayesian.MCMC.HMC (HMCConfig (..), HMCState (..), hmc)++-- | Configuration for MALA.+newtype MALAConfig = MALAConfig+ { malaStepSize :: Double+ }++-- | MALA is HMC with a single leapfrog step.+mala :: forall s d.+ (KnownShape s, KnownDType d)+ => (Tensor s d -> Builder (Tensor '[] d))+ -> Gradient s d+ -> MALAConfig+ -> Kernel s d (HMCState s d) (Info s d)+mala logpdf grad config =+ hmc logpdf grad (HMCConfig (malaStepSize config) 1)
+ src/HBayesian/MCMC/RandomWalk.hs view
@@ -0,0 +1,61 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module HBayesian.MCMC.RandomWalk+ ( RWConfig (..)+ , randomWalk+ ) where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HBayesian.Core+import HBayesian.HHLO.Ops+import qualified HBayesian.HHLO.RNG as RNG++newtype RWConfig = RWConfig+ { rwScale :: Double+ }++randomWalk :: forall s d.+ (KnownShape s, KnownDType d)+ => (Tensor s d -> Builder (Tensor '[] d))+ -> RWConfig+ -> SimpleKernel s d+randomWalk logpdf config = Kernel { kernelInit = kernelInit, kernelStep = kernelStep }+ where+ scaleVal = rwScale config++ kernelInit _key pos = do+ ld <- logpdf pos+ return $ State pos ld++ kernelStep key state = do+ (key1, key2) <- RNG.splitKey key+ let pos = statePosition state+ let ld = stateLogDensity state++ noise <- RNG.rngNormalF32 key1 >>= convert @s @'F32 @d+ scaleT <- constant @s @d scaleVal+ scaledNoise <- tmul noise scaleT+ pos' <- tadd pos scaledNoise++ ld' <- logpdf pos'++ diff <- tsub ld' ld+ zero <- constant @'[] @d 0.0+ logAlpha <- tminimum diff zero++ u <- RNG.rngUniformF32 key2 >>= convert @'[] @'F32 @d+ logU <- tlog u+ accept <- tlessThan logU logAlpha+ acceptS <- tbroadcast @'[] @s [] accept++ newPos <- tselect acceptS pos' pos+ newLd <- tselect accept ld' ld++ one <- constant @'[] @'I64 1+ infoAcceptProb <- texp logAlpha+ let info = Info infoAcceptProb accept one+ return (State newPos newLd, info)
+ src/HBayesian/PPL.hs view
@@ -0,0 +1,183 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++-- | A shallow probabilistic programming layer for HBayesian.+--+-- The PPL is a monad that carries:+-- * a read-only parameter vector of shape @[n]@+-- * a mutable log-density accumulator (scalar)+--+-- It desugars to the same 'Tensor s d -> Builder (Tensor '[] d)'+-- that samplers expect.+--+-- Example:+--+-- > myModel :: PPL 2 ()+-- > myModel = do+-- > alpha <- param 0+-- > beta <- param 1+-- > observe "alpha_prior" (normal 0.0 1.0) alpha+-- > observe "beta_prior" (normal 0.0 1.0) beta+-- > forM_ dataset $ \(x, y) -> do+-- > let mu = alpha + beta * x+-- > observe "y" (normal mu 0.5) y+-- >+-- > logpdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+-- > logpdf = runPPL myModel+module HBayesian.PPL+ ( -- * PPL monad+ PPL+ , runPPL+ , liftBuilder+ -- * Model vocabulary+ , param+ , observe+ -- * Distribution primitives+ , normal+ , normalT+ , uniform+ , uniformT+ , halfNormal+ , bernoulli+ , bernoulliT+ ) where++import GHC.TypeNats (Nat)++import HHLO.Core.Types+import HHLO.IR.Builder+import HBayesian.HHLO.Ops hiding (map)+import qualified HHLO.EDSL.Ops as EDSL++-----------------------------------------------------------------------------+-- PPL monad (manual implementation, no mtl dependency)+-----------------------------------------------------------------------------++-- | The PPL monad. Carries:+-- * a read-only parameter vector of shape @[n]@+-- * a mutable log-density accumulator (scalar)+newtype PPL (n :: Nat) a = PPL+ { unPPL :: Tensor '[n] 'F32 -> Tensor '[] 'F32 -> Builder (a, Tensor '[] 'F32)+ }++instance Functor (PPL n) where+ fmap f (PPL m) = PPL $ \theta acc -> do+ (x, acc') <- m theta acc+ return (f x, acc')++instance Applicative (PPL n) where+ pure x = PPL $ \_ acc -> return (x, acc)+ PPL mf <*> PPL mx = PPL $ \theta acc -> do+ (f, acc1) <- mf theta acc+ (x, acc2) <- mx theta acc1+ return (f x, acc2)++instance Monad (PPL n) where+ return = pure+ PPL mx >>= f = PPL $ \theta acc -> do+ (x, acc1) <- mx theta acc+ unPPL (f x) theta acc1++-- | Lift a raw 'Builder' action into the PPL.+liftBuilder :: Builder a -> PPL n a+liftBuilder mx = PPL $ \_ acc -> do+ x <- mx+ return (x, acc)++-- | Run a PPL model starting from a parameter vector and zero log-density.+-- Returns the accumulated log-posterior.+runPPL :: KnownShape '[n] => PPL n () -> Tensor '[n] 'F32 -> Builder (Tensor '[] 'F32)+runPPL (PPL m) theta = do+ zero <- tconstant 0.0+ ((), acc) <- m theta zero+ return acc++-----------------------------------------------------------------------------+-- Model vocabulary+-----------------------------------------------------------------------------++-- | Extract the i-th scalar parameter from the parameter vector.+param :: forall n. (KnownShape '[n], KnownDType 'F32)+ => Int -> PPL n (Tensor '[] 'F32)+param i = PPL $ \theta acc -> do+ x <- tslice1 @n @'F32 theta (fromIntegral i)+ return (x, acc)++-- | Condition the model on an observed value from a distribution.+-- The distribution is a function @value -> Builder (Tensor '[] 'F32)@+-- that computes the log-density.+observe :: String -> (Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)) -> Tensor '[] 'F32 -> PPL n ()+observe _name dist value = PPL $ \_ acc -> do+ ld <- dist value+ newAcc <- tadd acc ld+ return ((), newAcc)++-----------------------------------------------------------------------------+-- Distribution primitives+-----------------------------------------------------------------------------++-- | Normal distribution with constant mean and std (unnormalised log-pdf).+normal :: Double -> Double -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+normal mean std x = do+ meanT <- tconstant mean+ stdT <- tconstant std+ normalT meanT stdT x++-- | Normal distribution with tensor mean and std.+normalT :: Tensor '[] 'F32 -> Tensor '[] 'F32 -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+normalT meanT stdT x = do+ diff <- tsub x meanT+ ratio <- tdiv diff stdT+ sq <- tmul ratio ratio+ negHalf <- tconstant (-0.5)+ tmul negHalf sq++-- | Uniform distribution on @[a, b]@ with constant bounds.+uniform :: Double -> Double -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+uniform a b x = do+ aT <- tconstant a+ bT <- tconstant b+ uniformT aT bT x++-- | Uniform distribution with tensor bounds.+uniformT :: Tensor '[] 'F32 -> Tensor '[] 'F32 -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+uniformT aT bT x = do+ geA <- EDSL.compare x aT "GE"+ leB <- EDSL.compare x bT "LE"+ zero <- tconstant 0.0+ negInf <- tconstant (-1.0e30) -- proxy for -inf+ -- inside = geA && leB via nested tselect+ inner <- tselect leB zero negInf+ tselect geA inner negInf++-- | Half-normal distribution (positive support) with constant std.+halfNormal :: Double -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+halfNormal std x = do+ zero <- tconstant 0.0+ isPos <- EDSL.compare x zero "GE"+ stdT <- tconstant std+ lp <- normalT zero stdT x+ negInf <- tconstant (-1.0e30)+ tselect isPos lp negInf++-- | Bernoulli distribution with constant probability.+bernoulli :: Double -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+bernoulli p y = do+ pT <- tconstant p+ bernoulliT pT y++-- | Bernoulli distribution with tensor probability.+bernoulliT :: Tensor '[] 'F32 -> Tensor '[] 'F32 -> Builder (Tensor '[] 'F32)+bernoulliT pT y = do+ logP <- tlog pT+ one <- tconstant 1.0+ oneMinP <- tsub one pT+ logOneMinP <- tlog oneMinP+ term1 <- tmul y logP+ oneMinY <- tsub one y+ term2 <- tmul oneMinY logOneMinP+ tadd term1 term2
+ test/Main.hs view
@@ -0,0 +1,25 @@+module Main (main) where++import Test.Tasty+import qualified Test.Core as Core+import qualified Test.HHLO.Ops as Ops+import qualified Test.HHLO.RNG as RNG+import qualified Test.HHLO.Loops as Loops+import qualified Test.MCMC as MCMC+import qualified Test.Examples as Examples+import qualified Test.Chain as Chain+import qualified Test.PPL as PPL+import qualified Test.CorrelatedGaussian as CorrG++main :: IO ()+main = defaultMain $ testGroup "HBaysian Tests"+ [ Core.tests+ , Ops.tests+ , RNG.tests+ , Loops.tests+ , MCMC.tests+ , Examples.tests+ , Chain.tests+ , PPL.tests+ , CorrG.tests+ ]
+ test/Test/Chain.hs view
@@ -0,0 +1,33 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Test.Chain (tests) where++import Test.Tasty+import Test.Tasty.HUnit++import qualified LinearRegressionRandomWalk as Ex1+import qualified BivariateGaussianMALA as Ex4++import HBayesian.Chain+import HBayesian.Diagnostics++tests :: TestTree+tests = testGroup "Chain"+ [ testCase "LinearRegressionRandomWalk V2 runs and returns 10 samples" $ do+ (samples, diags) <- Ex1.runChainV2+ length samples @?= 10+ length diags @?= 10++ , testCase "LinearRegressionRandomWalk V2 acceptance rate is in [0,1]" $ do+ (_samples, diags) <- Ex1.runChainV2+ let rate = acceptanceRate diags+ assertBool "acceptance rate should be >= 0" (rate >= 0)+ assertBool "acceptance rate should be <= 1" (rate <= 1)++ , testCase "LinearRegressionRandomWalk V2 mean accept prob is in [0,1]" $ do+ (_samples, diags) <- Ex1.runChainV2+ let m = meanAcceptProb diags+ assertBool "mean accept prob should be >= 0" (m >= 0)+ assertBool "mean accept prob should be <= 1" (m <= 1)+ ]
+ test/Test/Core.hs view
@@ -0,0 +1,20 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Test.Core (tests) where++import Test.Tasty+import Test.Tasty.HUnit++import HBayesian.Core++-- | Core type sanity checks. Since Core is mostly type aliases and+-- data declarations, we verify that they can be constructed and that+-- the expected types line up.+tests :: TestTree+tests = testGroup "Core"+ [ testCase "Key newtype" $ do+ -- We can only test trivial construction here because Key wraps+ -- a Tensor which requires a Builder context.+ True @?= True+ ]
@@ -0,0 +1,213 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Test.CorrelatedGaussian (tests) where++import Control.Monad (forM_)+import Data.List (sort, transpose)+import Test.Tasty+import Test.Tasty.HUnit++import CorrelatedGaussianHMC+import HBayesian.Chain+import HBayesian.Diagnostics++-----------------------------------------------------------------------------+-- Helpers+-----------------------------------------------------------------------------++-- | Sample mean of a list.+mean :: [Float] -> Float+mean xs = sum xs / fromIntegral (length xs)++-- | Sample variance (unbiased).+variance :: [Float] -> Float+variance xs =+ let m = mean xs+ n = fromIntegral (length xs)+ in if n <= 1 then 0.0+ else sum [ (x - m) ^ 2 | x <- xs ] / (n - 1)++-- | Standard deviation.+stdDev :: [Float] -> Float+stdDev xs = sqrt (variance xs)++-- | Normal CDF approximation (Abramowitz & Stegun, formula 26.2.17).+-- Accurate to ~1e-7.+normalCdf :: Float -> Float+normalCdf x =+ let z = realToFrac x :: Double+ b1 = 0.319381530+ b2 = -0.356563782+ b3 = 1.781477937+ b4 = -1.821255978+ b5 = 1.330274429+ p = 0.2316419+ c = 0.39894228+ t = 1.0 / (1.0 + p * abs z)+ phi = c * exp (-0.5 * z * z)+ poly = t * (b1 + t * (b2 + t * (b3 + t * (b4 + t * b5))))+ y = 1.0 - phi * poly+ in realToFrac $ if z > 0 then y else 1.0 - y++-- | Empirical CDF at point x for sorted samples.+empiricalCdf :: Float -> [Float] -> Float+empiricalCdf x sortedXs =+ let n = fromIntegral (length sortedXs)+ count = fromIntegral (length (takeWhile (<= x) sortedXs))+ in count / n++-- | Kolmogorov-Smirnov statistic against a theoretical CDF.+ksStatistic :: (Float -> Float) -> [Float] -> Float+ksStatistic cdf xs =+ let sorted = sort xs+ n = fromIntegral (length sorted)+ diffs = [ max (abs (empiricalCdf x sorted - cdf x))+ (abs (empiricalCdfPrev x sorted - cdf x))+ | x <- sorted ]+ empiricalCdfPrev x sortedXs =+ let count = fromIntegral (length (takeWhile (< x) sortedXs))+ in count / n+ in maximum diffs++-- | Critical value for KS test at α = 0.01 for large n.+-- D_crit ≈ 1.628 / sqrt(n)+ksCriticalValue :: Int -> Float+ksCriticalValue n = 1.628 / sqrt (fromIntegral n)++-- | Mahalanobis distance squared: (x - μ)^T Λ (x - μ)+mahalanobisSq :: [Float] -> [Float] -> [[Float]] -> Float+mahalanobisSq x mu lam =+ let diff = zipWith (-) x mu+ -- For tridiagonal Λ, compute efficiently+ n = length diff+ in sum [ diff !! i * lam !! i !! j * diff !! j+ | i <- [0..n-1], j <- [0..n-1] ]++-- | Gelman-Rubin R-hat for a single dimension across multiple chains.+rHatSamples :: [[Float]] -> Float+rHatSamples chains =+ let m = fromIntegral (length chains) -- number of chains+ n = fromIntegral (length (head chains)) -- iterations per chain+ chainMeans = map mean chains+ chainVars = map variance chains+ w = mean chainVars -- within-chain variance+ b = n * variance chainMeans -- between-chain variance+ vHat = ((n - 1) / n) * w + (1 / n) * b -- pooled variance+ in if w <= 0 then 1.0 else sqrt (vHat / w)++-----------------------------------------------------------------------------+-- Tests+-----------------------------------------------------------------------------++-- | Number of chains for R-hat test.+numChains :: Int+numChains = 4++-- | Perturbation for dispersed initial values.+perturb :: [Float] -> [Float]+perturb = map (+ 0.5)++-- | Run a single chain and return samples.+runOneChain :: IO [Float]+runOneChain = do+ (samples, _diags) <- runChainV2+ return (map head samples) -- just first dim for some tests++tests :: TestTree+tests = testGroup "CorrelatedGaussian"+ [ testCase "HMC returns correct number of samples" testSampleCount+ , testCase "Marginal means are close to ground truth" testMarginalMeans+ , testCase "Marginal variances are close to 1.0" testMarginalVars+ , testCase "Marginal KS tests pass (dim 0)" testKS0+ , testCase "Marginal KS tests pass (dim 1)" testKS1+ , testCase "Marginal KS tests pass (dim 2)" testKS2+ , testCase "Marginal KS tests pass (dim 3)" testKS3+ , testCase "Marginal KS tests pass (dim 4)" testKS4+ , testCase "Mahalanobis distances have correct mean" testMahalanobisMean+ , testCase "R-hat < 1.1 across 4 chains" testRhat+ ]++-- | Test that we get the expected number of samples.+testSampleCount :: Assertion+testSampleCount = do+ (samples, _diags) <- runChainV2+ length samples @?= 2000++-- | Test that each marginal mean is within 3 SE of the true mean.+-- For AR(1) with uniform variance, marginal SD = 1, so SE = 1/sqrt(2000) ≈ 0.022.+testMarginalMeans :: Assertion+testMarginalMeans = do+ (samples, _diags) <- runChainV2+ let n = fromIntegral (length samples) :: Float+ se = 1.0 / sqrt n -- marginal SD = 1 for all dims+ thresh = 3.0 * se -- ~0.067+ forM_ (zip [0..4] targetMean) $ \(i, muTrue) -> do+ let xs = map (!! i) samples+ muHat = mean xs+ err = abs (muHat - muTrue)+ assertBool ("dim " ++ show i ++ " mean off by " ++ show err)+ (err < thresh)++-- | Test that each marginal variance is close to 1.0.+testMarginalVars :: Assertion+testMarginalVars = do+ (samples, _diags) <- runChainV2+ let thresh = 0.15 -- generous for 2000 samples+ forM_ [0..4] $ \i -> do+ let xs = map (!! i) samples+ varHat = variance xs+ err = abs (varHat - 1.0)+ assertBool ("dim " ++ show i ++ " variance off by " ++ show err)+ (err < thresh)++-- | Marginal KS tests for each dimension.+-- The marginal of an AR(1) Gaussian is N(mu_i, 1).+testKS :: Int -> Assertion+testKS i = do+ (samples, _diags) <- runChainV2+ let xs = map (!! i) samples+ mu = targetMean !! i+ cdf x = normalCdf ((x - mu) / 1.0)+ stat = ksStatistic cdf xs+ crit = ksCriticalValue (length xs)+ assertBool ("dim " ++ show i ++ " KS stat " ++ show stat ++ " > crit " ++ show crit)+ (stat < crit)++testKS0, testKS1, testKS2, testKS3, testKS4 :: Assertion+testKS0 = testKS 0+testKS1 = testKS 1+testKS2 = testKS 2+testKS3 = testKS 3+testKS4 = testKS 4++-- | Test that mean of Mahalanobis distances ≈ k = 5.+-- E[d²] = k for χ²_k.+testMahalanobisMean :: Assertion+testMahalanobisMean = do+ (samples, _diags) <- runChainV2+ let dists = [ mahalanobisSq s targetMean targetPrecision | s <- samples ]+ meanD = mean dists+ -- For χ²_5, E[d²] = 5, Var[d²] = 10+ -- SE of mean = sqrt(10/n) ≈ 0.071+ se = sqrt (10.0 / fromIntegral (length samples))+ thresh = 3.0 * se -- ~0.21+ err = abs (meanD - 5.0)+ assertBool ("Mahalanobis mean " ++ show meanD ++ " off by " ++ show err)+ (err < thresh)++-- | Test Gelman-Rubin across 4 chains.+testRhat :: Assertion+testRhat = do+ let config = defaultChainConfig { ccNumIterations = 500 }+ ck = compileHMC (makeKernel (HMCConfig 0.1 10))+ gaussianLogPdf gaussianGrad+ results <- parallelChains numChains perturb ck+ (replicate targetDim 0.0) config+ let chains = map fst results+ -- Transpose: get per-dimension chains+ perDimChains = [ [ map (!! i) c | c <- chains ] | i <- [0..4] ]+ rhats = map rHatSamples perDimChains+ forM_ (zip [0..4] rhats) $ \(i, r) -> do+ assertBool ("dim " ++ show i ++ " R-hat = " ++ show r ++ " >= 1.1")+ (r < 1.1)
+ test/Test/Examples.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Examples (tests) where++import Data.Text (Text)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import qualified LinearRegressionRandomWalk as Ex1+import qualified GaussianProcessEllipticalSlice as Ex2+import qualified LogisticRegressionHMC as Ex3+import qualified BivariateGaussianMALA as Ex4++tests :: TestTree+tests = testGroup "Examples"+ [ testCase "LinearRegressionRandomWalk renders MLIR" $ do+ let mlir = Ex1.renderStepMlir+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains multiply" (T.isInfixOf "stablehlo.multiply" mlir)+ assertBool "contains compare" (T.isInfixOf "stablehlo.compare" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "GaussianProcessEllipticalSlice renders MLIR" $ do+ let mlir = Ex2.renderStepMlir+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains cosine" (T.isInfixOf "stablehlo.cosine" mlir)+ assertBool "contains sine" (T.isInfixOf "stablehlo.sine" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "LogisticRegressionHMC renders MLIR" $ do+ let mlir = Ex3.renderStepMlir+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains multiply" (T.isInfixOf "stablehlo.multiply" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "BivariateGaussianMALA renders MLIR" $ do+ let mlir = Ex4.renderStepMlir+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains multiply" (T.isInfixOf "stablehlo.multiply" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)+ ]
+ test/Test/HHLO/Loops.hs view
@@ -0,0 +1,69 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.HHLO.Loops (tests) where++import Data.Text (Text)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..), Module(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HBayesian.HHLO.Ops+import HBayesian.HHLO.Loops++-- | Build a 3-tuple and render it.+render3T :: forall s1 d1 s2 d2 s3 d3.+ (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+ => [FuncArg] -> Builder (Tuple '[s1, s2, s3] '[d1, d2, d3]) -> Text+render3T args b = render $ moduleFromBuilderT @'[s1, s2, s3] @'[d1, d2, d3] "main" args b++tests :: TestTree+tests = testGroup "HHLO.Loops"+ [ testCase "whileLoop3 renders stablehlo.while" $ do+ let mlir = render3T @'[2] @'F32 @'[2] @'F32 @'[] @'I64+ [ FuncArg "a" (TensorType [2] F32)+ , FuncArg "b" (TensorType [2] F32)+ , FuncArg "c" (TensorType [] I64)+ ] $ do+ a0 <- arg @'[2] @'F32+ b0 <- arg @'[2] @'F32+ c0 <- arg @'[] @'I64+ (a1, b1, c1) <- whileLoop3 a0 b0 c0+ (\a b c -> do+ limit <- constant @'[] @'I64 10+ lessThan c limit)+ (\a b c -> do+ a' <- tadd a b+ b' <- tmul b a+ one <- constant @'[] @'I64 1+ c' <- add c one+ return (a', b', c'))+ return $ a1 ::: b1 ::: c1 ::: TNil+ assertBool "contains stablehlo.while" (T.isInfixOf "stablehlo.while" mlir)+ assertBool "contains stablehlo.return" (T.isInfixOf "stablehlo.return" mlir)++ , testCase "conditional3 renders stablehlo.if" $ do+ let mlir = render3T @'[2] @'F32 @'[2] @'F32 @'[] @'I64+ [ FuncArg "p" (TensorType [] Bool)+ ] $ do+ p <- arg @'[] @'Bool+ (a, b, c) <- conditional3 p+ (do+ x <- constant @'[2] @'F32 1.0+ y <- constant @'[2] @'F32 2.0+ z <- constant @'[] @'I64 3+ return (x, y, z))+ (do+ x <- constant @'[2] @'F32 0.0+ y <- constant @'[2] @'F32 0.0+ z <- constant @'[] @'I64 0+ return (x, y, z))+ return $ a ::: b ::: c ::: TNil+ assertBool "contains stablehlo.if" (T.isInfixOf "stablehlo.if" mlir)+ ]
+ test/Test/HHLO/Ops.hs view
@@ -0,0 +1,70 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.HHLO.Ops (tests) where++import Data.Text (Text)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HBayesian.HHLO.Ops++-- | Render a single-result builder to text.+render1 :: forall s d. (KnownShape s, KnownDType d)+ => [FuncArg] -> Builder (Tensor s d) -> Text+render1 args b = render $ moduleFromBuilder @s @d "main" args b++tests :: TestTree+tests = testGroup "HHLO.Ops"+ [ testCase "sqrt op renders" $ do+ let mlir = render1 @'[2,2] @'F32+ [ FuncArg "x" (TensorType [2,2] F32) ] $ do+ x <- arg @'[2,2] @'F32+ sqrt' x+ let expected = "module {\n func.func @main(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {\n %0 = stablehlo.sqrt %arg0 : (tensor<2x2xf32>) -> tensor<2x2xf32>\n return %0 : tensor<2x2xf32>\n }\n}"+ mlir @?= expected++ , testCase "sin op renders" $ do+ let mlir = render1 @'[3] @'F64+ [ FuncArg "x" (TensorType [3] F64) ] $ do+ x <- arg @'[3] @'F64+ sin' x+ let expected = "module {\n func.func @main(%arg0: tensor<3xf64>) -> tensor<3xf64> {\n %0 = stablehlo.sine %arg0 : (tensor<3xf64>) -> tensor<3xf64>\n return %0 : tensor<3xf64>\n }\n}"+ mlir @?= expected++ , testCase "cos op renders" $ do+ let mlir = render1 @'[3] @'F64+ [ FuncArg "x" (TensorType [3] F64) ] $ do+ x <- arg @'[3] @'F64+ cos' x+ let expected = "module {\n func.func @main(%arg0: tensor<3xf64>) -> tensor<3xf64> {\n %0 = stablehlo.cosine %arg0 : (tensor<3xf64>) -> tensor<3xf64>\n return %0 : tensor<3xf64>\n }\n}"+ mlir @?= expected++ , testCase "pow op renders" $ do+ let mlir = render1 @'[2] @'F32+ [ FuncArg "x" (TensorType [2] F32)+ , FuncArg "y" (TensorType [2] F32)+ ] $ do+ x <- arg @'[2] @'F32+ y <- arg @'[2] @'F32+ pow' x y+ let expected = "module {\n func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {\n %0 = stablehlo.power %arg0, %arg1 : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>\n return %0 : tensor<2xf32>\n }\n}"+ mlir @?= expected++ , testCase "element-wise lessThan renders" $ do+ let mlir = render1 @'[2] @'Bool+ [ FuncArg "x" (TensorType [2] F32)+ , FuncArg "y" (TensorType [2] F32)+ ] $ do+ x <- arg @'[2] @'F32+ y <- arg @'[2] @'F32+ lessThanEW x y+ let expected = "module {\n func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xi1> {\n %0 = \"stablehlo.compare\"(%arg0, %arg1) {comparison_direction = \"LT\"} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>\n return %0 : tensor<2xi1>\n }\n}"+ mlir @?= expected+ ]
+ test/Test/HHLO/RNG.hs view
@@ -0,0 +1,58 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.HHLO.RNG (tests) where++import Data.Text (Text)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HBayesian.HHLO.Ops+import HBayesian.HHLO.RNG++render1 :: forall s d. (KnownShape s, KnownDType d)+ => [FuncArg] -> Builder (Tensor s d) -> Text+render1 args b = render $ moduleFromBuilder @s @d "main" args b++render2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+ => [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Text+render2 args b = render $ moduleFromBuilder2 @s1 @d1 @s2 @d2 "main" args b++tests :: TestTree+tests = testGroup "HHLO.RNG"+ [ testCase "splitKey renders two rng_bit_generator ops" $ do+ let mlir = render2 @'[2] @'UI64 @'[2] @'UI64+ [ FuncArg "key" (TensorType [2] UI64) ] $ do+ k <- arg @'[2] @'UI64+ (k1, k2) <- splitKey (Key k)+ returnTuple2 (unKey k1) (unKey k2)+ -- We just check that the expected ops appear+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains THREE_FRY" (T.isInfixOf "THREE_FRY" mlir)++ , testCase "rngUniformF32 renders" $ do+ let mlir = render1 @'[3] @'F32+ [ FuncArg "key" (TensorType [2] UI64) ] $ do+ k <- arg @'[2] @'UI64+ rngUniformF32 (Key k)+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains convert" (T.isInfixOf "stablehlo.convert" mlir)+ assertBool "contains divide" (T.isInfixOf "stablehlo.divide" mlir)++ , testCase "rngNormalF32 renders" $ do+ let mlir = render1 @'[3] @'F32+ [ FuncArg "key" (TensorType [2] UI64) ] $ do+ k <- arg @'[2] @'UI64+ rngNormalF32 (Key k)+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains log" (T.isInfixOf "stablehlo.log" mlir)+ assertBool "contains sqrt" (T.isInfixOf "stablehlo.sqrt" mlir)+ assertBool "contains cosine" (T.isInfixOf "stablehlo.cosine" mlir)+ ]
+ test/Test/MCMC.hs view
@@ -0,0 +1,150 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.MCMC (tests) where++import Data.Text (Text)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)+import HBayesian.Core+import HBayesian.HHLO.Ops+import qualified HBayesian.HHLO.RNG as RNG+import HBayesian.MCMC.EllipticalSlice+import HBayesian.MCMC.HMC+import HBayesian.MCMC.MALA++render1 :: forall s d. (KnownShape s, KnownDType d)+ => [FuncArg] -> Builder (Tensor s d) -> Text+render1 args b = render $ moduleFromBuilder @s @d "main" args b++-- | Dummy log-density for a 1-D standard normal.+-- log p(x) = -0.5 * x^2 (ignoring constant)+stdNormalLogPdf :: Tensor '[1] 'F32 -> Builder (Tensor '[] 'F32)+stdNormalLogPdf x = do+ xSq <- tmul x x+ negHalf <- constant @'[1] @'F32 (-0.5)+ negHalfXSq <- tmul negHalf xSq+ tsumAll negHalfXSq++-- | Dummy gradient for the 1-D standard normal.+-- d/dx (-0.5 * x^2) = -x+stdNormalGrad :: Gradient '[1] 'F32+stdNormalGrad x = do+ negOne <- constant @'[1] @'F32 (-1.0)+ tmul negOne x++tests :: TestTree+tests = testGroup "MCMC"+ [ testCase "RandomWalk kernelStep renders" $ do+ let k = Kernel+ { kernelInit = \_key pos -> do+ ld <- stdNormalLogPdf pos+ return (State pos ld)+ , kernelStep = \key state -> do+ (k1, k2) <- RNG.splitKey key+ let pos = statePosition state+ let ld = stateLogDensity state+ noise <- RNG.rngNormalF32 k1 >>= convert @'[1] @'F32 @'F32+ scaleT <- constant @'[1] @'F32 0.1+ scaledNoise <- tmul noise scaleT+ pos' <- tadd pos scaledNoise+ ld' <- stdNormalLogPdf pos'+ diff <- tsub ld' ld+ zero <- constant @'[] @'F32 0.0+ logAlpha <- tminimum diff zero+ u <- RNG.rngUniformF32 k2 >>= convert @'[] @'F32 @'F32+ logU <- tlog u+ accept <- tlessThan logU logAlpha+ acceptS <- tbroadcast @'[] @'[1] [] accept+ newPos <- tselect acceptS pos' pos+ newLd <- tselect accept ld' ld+ one <- constant @'[] @'I64 1+ acceptProb <- texp logAlpha+ let info = Info acceptProb accept one+ return (State newPos newLd, info)+ }+ let mlir = render1 @'[1] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [1] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[1] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep k (Key key) (State pos ld)+ return (statePosition state')+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains subtract" (T.isInfixOf "stablehlo.subtract" mlir)+ assertBool "contains compare" (T.isInfixOf "stablehlo.compare" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "EllipticalSlice kernelStep renders" $ do+ let k = ellipticalSlice stdNormalLogPdf+ let mlir = render1 @'[1] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [1] F32)+ , FuncArg "ld" (TensorType [] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[1] @'F32+ ld <- arg @'[] @'F32+ (state', _info) <- kernelStep k (Key key) (State pos ld)+ return (statePosition state')+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains cosine" (T.isInfixOf "stablehlo.cosine" mlir)+ assertBool "contains sine" (T.isInfixOf "stablehlo.sine" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "HMC kernelStep renders" $ do+ let config = HMCConfig { hmcStepSize = 0.1, hmcNumLeapfrogSteps = 2 }+ let k = hmc stdNormalLogPdf stdNormalGrad config+ let mlir = render1 @'[1] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [1] F32)+ , FuncArg "p" (TensorType [1] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [1] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[1] @'F32+ p <- arg @'[1] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[1] @'F32+ (state', _info) <- kernelStep k (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains multiply" (T.isInfixOf "stablehlo.multiply" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)++ , testCase "MALA kernelStep renders" $ do+ let config = MALAConfig { malaStepSize = 0.1 }+ let k = mala stdNormalLogPdf stdNormalGrad config+ let mlir = render1 @'[1] @'F32+ [ FuncArg "key" (TensorType [2] UI64)+ , FuncArg "pos" (TensorType [1] F32)+ , FuncArg "p" (TensorType [1] F32)+ , FuncArg "ld" (TensorType [] F32)+ , FuncArg "g" (TensorType [1] F32)+ ] $ do+ key <- arg @'[2] @'UI64+ pos <- arg @'[1] @'F32+ p <- arg @'[1] @'F32+ ld <- arg @'[] @'F32+ g <- arg @'[1] @'F32+ (state', _info) <- kernelStep k (Key key) (HMCState pos p ld g)+ return (hmcPosition state')+ assertBool "contains rng_bit_generator" (T.isInfixOf "stablehlo.rng_bit_generator" mlir)+ assertBool "contains add" (T.isInfixOf "stablehlo.add" mlir)+ assertBool "contains multiply" (T.isInfixOf "stablehlo.multiply" mlir)+ assertBool "contains select" (T.isInfixOf "stablehlo.select" mlir)+ ]
+ test/Test/PPL.hs view
@@ -0,0 +1,56 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.PPL (tests) where++import Data.List (isInfixOf)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty (render)++import HBayesian.Core+import HBayesian.HHLO.Ops hiding (map)+import HBayesian.PPL+import HBayesian.MCMC.RandomWalk+import HBayesian.Chain++-- | A tiny PPL model: two parameters with normal priors and one observation.+tinyModel :: PPL 2 ()+tinyModel = do+ alpha <- param 0+ beta <- param 1+ observe "alpha_prior" (normal 0.0 1.0) alpha+ observe "beta_prior" (normal 0.0 1.0) beta+ -- observation: y = 1.0 at x = 0.5, likelihood N(alpha + beta * 0.5, 0.5)+ x <- liftBuilder $ tconstant 0.5+ mu <- liftBuilder $ tadd alpha =<< tmul beta x+ y <- liftBuilder $ tconstant 1.0+ sigma <- liftBuilder $ tconstant 0.5+ observe "y" (normalT mu sigma) y++-- | The same model written manually as a log-posterior.+tinyLogPdf :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+tinyLogPdf = runPPL tinyModel++-- | Render the PPL-derived log-posterior to MLIR for inspection.+tinyMlir :: String+tinyMlir = show $ render $ moduleFromBuilder @'[] @'F32 "main"+ [FuncArg "theta" (TensorType [2] F32)] $ do+ theta <- arg @'[2] @'F32+ tinyLogPdf theta++tests :: TestTree+tests = testGroup "PPL"+ [ testCase "PPL-derived model renders non-empty MLIR" $ do+ assertBool "MLIR should be non-empty" (not (null tinyMlir))++ , testCase "PPL-derived model contains expected ops" $ do+ assertBool "should contain stablehlo.constant" ("stablehlo.constant" `isInfixOf` tinyMlir)+ assertBool "should contain stablehlo.add" ("stablehlo.add" `isInfixOf` tinyMlir)+ assertBool "should contain stablehlo.multiply" ("stablehlo.multiply" `isInfixOf` tinyMlir)+ ]