UncertainTea - shohei81.dev

UncertainTea は、Gen 風のフロントエンドと静的実行モデルを持つ実験的な Julia 確率的プログラミングパッケージです。logjoint、バッチ化されたチェーン・パーティクル、パラメータ変換が密なレイアウトとバックエンド向きの制御フローで動くよう、モデル構造を静的に保つことだけに注力しています。

できること

@tea / @tea (static) による Gen 風のモデリング。チルダ構文、明示的なアドレス、階層アドレス、choicemap による外部条件付けに対応
modelspec・parameterlayout・executionplan による静的モデルの内省と、backend_report によるバックエンドの支援判定
generate・assess・logjoint・非拘束変換、そしてバッチ化された logjoint / 勾配 API
推論: hmc・nuts・batched_chees（ChEES-HMC）・batched_meads（MEADS）・gibbs・batched_advi・batched_svgd・batched_smc・nested_sampling など
推定と診断: pathfinder・map_estimate・laplace_approximation、sbc（シミュレーションによるキャリブレーション）・waic・psis_loo / loo・事後予測 predict
エコシステム連携: to_mcmcchains（MCMCChains / StatsPlots）、to_arviz_dict（Python ArviZ）、Tables.jl インターフェース
GPU 向けの実験的なバックエンド: KernelAbstractions によるデバイス密度 API（device_batched_logjoint ほか）と Metal エクステンション

技術構成

言語: Julia 1.10+（Apache 2.0）
設計: 単一の制約「モデル構造を静的かつ密に保つ」。CPU リファレンス実装を先に、デバイスバックエンド（KernelAbstractions + Metal）を並行開発中
評価: CI とベンチマーク、test/ による回帰テスト

Turing との互換性や動的トレースの完全サポートには意図的に注力していません。Gen 風の表面構文 → 静的セマンティクス → 密なパラメータレイアウト → CPU リファレンス → GPU バックエンド、という順で進んでいます。

ドキュメント

ドキュメントサイト — はじめ方・推論ガイド・実行可能な例・API リファレンス
research notes・architecture・batched inference

試してみる

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)

現在 0.2.0 の実験リリースです。API やモデルの制約は、静的 IR とバックエンド契約の収束に合わせて変わる可能性があります。