UncertainTea

UncertainTea is an experimental probabilistic programming package for Julia with a Gen-like frontend and a static execution model. The project is built around one constraint: keep model structure static enough that logjoint, batched chains or particles, and parameter transforms can run over dense layouts and backend-friendly control flow.

What it provides

• Gen-like modeling with @tea / @tea (static): tilde syntax, explicit and hierarchical addresses, external conditioning via choicemap
• Static model introspection through modelspec, parameterlayout, executionplan, and backend reports
• CPU reference evaluation with generate, assess, logjoint, unconstrained transforms, and batched logjoint/gradient APIs
• Inference: hmc, nuts, batched_chees (ChEES-HMC), batched_meads (MEADS), gibbs, batched_advi, batched_svgd, batched_smc, nested_sampling, and more
• Estimation and diagnostics: pathfinder, map_estimate, laplace_approximation, sbc (simulation-based calibration), waic, psis_loo / loo, and posterior-predictive predict
• Ecosystem interop via package extensions: to_mcmcchains (MCMCChains / StatsPlots), to_arviz_dict (Python ArviZ), and the Tables.jl interface
• Experimental GPU-oriented lowering: backend_report support checks plus a KernelAbstractions device backend (device_batched_logjoint, device_batched_logjoint_gradient) with a Metal extension

Design

• Language: Julia 1.10+ (Apache 2.0)
• Approach: a single constraint, static, dense model structure. CPU reference first, then a device backend (KernelAbstractions + Metal) developed in parallel
• Quality: CI, benchmarks, and a test/ regression suite

UncertainTea is intentionally not centered on Turing compatibility or unrestricted dynamic traces. The path is: Gen-like surface syntax → static semantics → dense parameter layouts → CPU reference → GPU backend.

Documentation

• Documentation site: getting started, inference guide, executable examples, API reference
• research notes · architecture · batched inference

Try it

using Pkg
Pkg.add(url = "https://github.com/shohei81/UncertainTea.jl.git")

using UncertainTea             # @tea DSL, choicemap, distributions
using UncertainTea.Inference   # hmc_chains, nuts_chains, ...

@tea (static) function gaussian_mean()
    mu ~ normal(0.0f0, 1.0f0)
    {:y} ~ normal(mu, 1.0f0)
    return mu
end

constraints = choicemap((:y, 0.3f0))
chains = hmc_chains(gaussian_mean, (), constraints; num_chains=4, num_samples=100, num_warmup=100)

UncertainTea 0.2.0 is an experimental release. APIs and model restrictions may change as the static IR and backend contract converge.