oxilean_std/probabilistic_programming/
functions.rs

1//! Auto-generated module
2//!
3//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
4
5use oxilean_kernel::Node;
6use oxilean_kernel::{BinderInfo, Declaration, Environment, Expr, Level, Name};
7
8use super::types::{Distribution, ParticleFilter};
9
10pub fn app(f: Expr, a: Expr) -> Expr {
11    Expr::App(Node::new(f), Node::new(a))
12}
13pub fn app2(f: Expr, a: Expr, b: Expr) -> Expr {
14    app(app(f, a), b)
15}
16pub fn app3(f: Expr, a: Expr, b: Expr, c: Expr) -> Expr {
17    app(app2(f, a, b), c)
18}
19pub fn cst(s: &str) -> Expr {
20    Expr::Const(Name::str(s), vec![])
21}
22pub fn prop() -> Expr {
23    Expr::Sort(Level::zero())
24}
25pub fn type0() -> Expr {
26    Expr::Sort(Level::succ(Level::zero()))
27}
28pub fn type1() -> Expr {
29    Expr::Sort(Level::succ(Level::succ(Level::zero())))
30}
31pub fn pi(bi: BinderInfo, name: &str, dom: Expr, body: Expr) -> Expr {
32    Expr::Pi(bi, Name::str(name), Node::new(dom), Node::new(body))
33}
34pub fn arrow(a: Expr, b: Expr) -> Expr {
35    pi(BinderInfo::Default, "_", a, b)
36}
37pub fn bvar(n: u32) -> Expr {
38    Expr::BVar(n)
39}
40pub fn nat_ty() -> Expr {
41    cst("Nat")
42}
43pub fn real_ty() -> Expr {
44    cst("Real")
45}
46pub fn bool_ty() -> Expr {
47    cst("Bool")
48}
49pub fn list_ty(elem: Expr) -> Expr {
50    app(cst("List"), elem)
51}
52/// `Measure : Type → Type` — a σ-finite measure on a measurable space.
53pub fn measure_ty() -> Expr {
54    arrow(type0(), type0())
55}
56/// `SigmaAlgebra : Type → Type` — a σ-algebra of subsets.
57pub fn sigma_algebra_ty() -> Expr {
58    arrow(type0(), type0())
59}
60/// `MeasurableSpace : Type` — a type equipped with a σ-algebra.
61pub fn measurable_space_ty() -> Expr {
62    type0()
63}
64/// `ProbabilityMonad : (Type → Type)` — the distribution / Giry monad.
65pub fn probability_monad_ty() -> Expr {
66    arrow(type0(), type0())
67}
68/// `Kernel : Type → Type → Type` — a Markov kernel k(x, A).
69pub fn kernel_ty() -> Expr {
70    arrow(type0(), arrow(type0(), type0()))
71}
72/// `Sampler : Type → Type` — a procedure that draws samples from a distribution.
73pub fn sampler_ty() -> Expr {
74    arrow(type0(), type0())
75}
76/// `Density : Type → Type` — a density/pmf function on a type.
77pub fn density_ty() -> Expr {
78    arrow(type0(), arrow(type0(), real_ty()))
79}
80/// `PPLProgram : Type → Type` — a probabilistic program returning values of type A.
81pub fn ppl_program_ty() -> Expr {
82    arrow(type0(), type0())
83}
84/// `ELBO : (Type → Type) → Type → Real` — evidence lower bound for variational inference.
85pub fn elbo_ty() -> Expr {
86    arrow(arrow(type0(), type0()), arrow(type0(), real_ty()))
87}
88/// `ImportanceWeight : Type → Real` — self-normalised importance weight.
89pub fn importance_weight_ty() -> Expr {
90    arrow(type0(), real_ty())
91}
92/// `ParticleFilter : Type → Type` — sequential Monte Carlo state estimator.
93pub fn particle_filter_ty() -> Expr {
94    arrow(type0(), type0())
95}
96/// `GradientEstimator : Type → Type` — a Monte Carlo gradient estimator.
97pub fn gradient_estimator_ty() -> Expr {
98    arrow(type0(), type0())
99}
100/// **Measure-Theoretic Bayes**: the posterior is the Radon-Nikodym derivative
101/// of the joint w.r.t. the marginal likelihood.
102///
103/// `bayes_measure_theory : ∀ (prior : Measure X) (likelihood : Kernel X Y),
104///   Posterior prior likelihood = RadonNikodym (Joint prior likelihood) (Marginal likelihood)`
105pub fn bayes_measure_theory_ty() -> Expr {
106    pi(
107        BinderInfo::Default,
108        "X",
109        type0(),
110        pi(
111            BinderInfo::Default,
112            "Y",
113            type0(),
114            pi(
115                BinderInfo::Default,
116                "prior",
117                app(cst("Measure"), bvar(1)),
118                pi(
119                    BinderInfo::Default,
120                    "likelihood",
121                    app2(cst("Kernel"), bvar(2), bvar(1)),
122                    app2(
123                        cst("Eq"),
124                        app2(cst("Posterior"), bvar(1), bvar(0)),
125                        app2(
126                            cst("RadonNikodym"),
127                            app2(cst("Joint"), bvar(1), bvar(0)),
128                            app(cst("Marginal"), bvar(0)),
129                        ),
130                    ),
131                ),
132            ),
133        ),
134    )
135}
136/// **Giry Monad Laws**: the distribution monad satisfies monad laws.
137///
138/// `giry_monad_laws : MonadLaws ProbabilityMonad`
139pub fn giry_monad_laws_ty() -> Expr {
140    app(cst("MonadLaws"), cst("ProbabilityMonad"))
141}
142/// **Importance Sampling Consistency**: the IS estimator is consistent as N→∞.
143///
144/// `is_consistency : ∀ (f : X → Real) (q p : Measure X),
145///   AbsContinuous p q → ConsistentEstimator (ISSampler f q p)`
146pub fn is_consistency_ty() -> Expr {
147    pi(
148        BinderInfo::Default,
149        "X",
150        type0(),
151        pi(
152            BinderInfo::Default,
153            "f",
154            arrow(bvar(0), real_ty()),
155            pi(
156                BinderInfo::Default,
157                "q",
158                app(cst("Measure"), bvar(1)),
159                pi(
160                    BinderInfo::Default,
161                    "p",
162                    app(cst("Measure"), bvar(2)),
163                    arrow(
164                        app2(cst("AbsContinuous"), bvar(0), bvar(1)),
165                        app(
166                            cst("ConsistentEstimator"),
167                            app3(cst("ISSampler"), bvar(3), bvar(1), bvar(0)),
168                        ),
169                    ),
170                ),
171            ),
172        ),
173    )
174}
175/// **ELBO Lower Bound**: ELBO(q) ≤ log p(x) for all variational families q.
176///
177/// `elbo_lower_bound : ∀ (q : Measure Z) (p : Joint Z X) (x : X),
178///   ELBO q p x ≤ LogMarginalLikelihood p x`
179pub fn elbo_lower_bound_ty() -> Expr {
180    pi(
181        BinderInfo::Default,
182        "Z",
183        type0(),
184        pi(
185            BinderInfo::Default,
186            "X",
187            type0(),
188            pi(
189                BinderInfo::Default,
190                "q",
191                app(cst("Measure"), bvar(1)),
192                pi(
193                    BinderInfo::Default,
194                    "p",
195                    app2(cst("JointMeasure"), bvar(2), bvar(1)),
196                    pi(
197                        BinderInfo::Default,
198                        "x",
199                        bvar(2),
200                        app2(
201                            cst("Real.le"),
202                            app3(cst("ELBO"), bvar(2), bvar(1), bvar(0)),
203                            app2(cst("LogMarginalLikelihood"), bvar(1), bvar(0)),
204                        ),
205                    ),
206                ),
207            ),
208        ),
209    )
210}
211/// **Reparameterisation Gradient**: the reparameterised gradient estimator is
212/// an unbiased estimator of ∇_φ E_{z~q_φ}[f(z)].
213///
214/// `reparam_unbiased : ∀ (φ : Params) (f : Z → Real),
215///   Unbiased (ReparamGradient f) (GradExpectation f φ)`
216pub fn reparam_unbiased_ty() -> Expr {
217    pi(
218        BinderInfo::Default,
219        "Z",
220        type0(),
221        pi(
222            BinderInfo::Default,
223            "Params",
224            type0(),
225            pi(
226                BinderInfo::Default,
227                "phi",
228                bvar(0),
229                pi(
230                    BinderInfo::Default,
231                    "f",
232                    arrow(bvar(2), real_ty()),
233                    app2(
234                        cst("Unbiased"),
235                        app(cst("ReparamGradient"), bvar(0)),
236                        app2(cst("GradExpectation"), bvar(0), bvar(1)),
237                    ),
238                ),
239            ),
240        ),
241    )
242}
243/// **HMC Correctness**: Hamiltonian Monte Carlo leaves the target distribution invariant.
244///
245/// `hmc_invariant : ∀ (target : Measure X), InvariantUnder (HMCKernel target) target`
246pub fn hmc_invariant_ty() -> Expr {
247    pi(
248        BinderInfo::Default,
249        "X",
250        type0(),
251        pi(
252            BinderInfo::Default,
253            "target",
254            app(cst("Measure"), bvar(0)),
255            app2(
256                cst("InvariantUnder"),
257                app(cst("HMCKernel"), bvar(0)),
258                bvar(0),
259            ),
260        ),
261    )
262}
263/// **SMC Consistency**: sequential Monte Carlo converges to the true filtering distribution.
264pub fn smc_consistency_ty() -> Expr {
265    pi(
266        BinderInfo::Default,
267        "X",
268        type0(),
269        pi(
270            BinderInfo::Default,
271            "ssm",
272            app(cst("StateSpaceModel"), bvar(0)),
273            app(cst("ConsistentFilteringDist"), bvar(0)),
274        ),
275    )
276}
277/// **Stochastic Variational Inference (SVI) Convergence**: SVI converges to a local ELBO maximum.
278///
279/// `svi_convergence : ∀ (q_family : VariationalFamily) (lr_schedule : LRSchedule),
280///   ConvergesToLocalOptimum (SVIOptimizer q_family lr_schedule)`
281pub fn svi_convergence_ty() -> Expr {
282    pi(
283        BinderInfo::Default,
284        "VF",
285        type0(),
286        pi(
287            BinderInfo::Default,
288            "lr",
289            type0(),
290            app(
291                cst("ConvergesToLocalOptimum"),
292                app2(cst("SVIOptimizer"), bvar(1), bvar(0)),
293            ),
294        ),
295    )
296}
297/// **Normalizing Flow Change of Variables**: the pushforward density satisfies the
298/// change-of-variables formula.
299///
300/// `normalizing_flow_cov : ∀ (f : X → Y) (p : Measure X),
301///   Bijective f → Density (Pushforward f p) y = Density p (f_inv y) * AbsDetJac f_inv y`
302pub fn normalizing_flow_cov_ty() -> Expr {
303    pi(
304        BinderInfo::Default,
305        "X",
306        type0(),
307        pi(
308            BinderInfo::Default,
309            "Y",
310            type0(),
311            pi(
312                BinderInfo::Default,
313                "f",
314                arrow(bvar(1), bvar(0)),
315                pi(
316                    BinderInfo::Default,
317                    "p",
318                    app(cst("Measure"), bvar(2)),
319                    arrow(
320                        app(cst("Bijective"), bvar(1)),
321                        app(cst("FlowDensityEq"), bvar(2)),
322                    ),
323                ),
324            ),
325        ),
326    )
327}
328/// **Score Function Estimator Unbiasedness**: the REINFORCE estimator is an unbiased
329/// gradient estimator under mild regularity conditions.
330///
331/// `score_fn_unbiased : ∀ (q_phi : ParametricMeasure) (f : Z → Real),
332///   RegularFamily q_phi → Unbiased (ScoreFnGrad q_phi f) (GradELBO q_phi f)`
333pub fn score_fn_unbiased_ty() -> Expr {
334    pi(
335        BinderInfo::Default,
336        "Z",
337        type0(),
338        pi(
339            BinderInfo::Default,
340            "q_phi",
341            app(cst("ParametricMeasure"), bvar(0)),
342            pi(
343                BinderInfo::Default,
344                "f",
345                arrow(bvar(1), real_ty()),
346                arrow(
347                    app(cst("RegularFamily"), bvar(1)),
348                    app2(
349                        cst("Unbiased"),
350                        app2(cst("ScoreFnGrad"), bvar(1), bvar(0)),
351                        app2(cst("GradELBO"), bvar(1), bvar(0)),
352                    ),
353                ),
354            ),
355        ),
356    )
357}
358/// **Pathwise Gradient Unbiasedness**: the reparameterised (pathwise) gradient is
359/// an unbiased estimator when the reparameterisation is differentiable.
360///
361/// `pathwise_gradient_unbiased : ∀ (g : Eps → Phi → Z) (f : Z → Real),
362///   DiffReparameterisation g → Unbiased (PathwiseGrad g f) (GradELBO (Reparam g) f)`
363pub fn pathwise_gradient_unbiased_ty() -> Expr {
364    pi(
365        BinderInfo::Default,
366        "Eps",
367        type0(),
368        pi(
369            BinderInfo::Default,
370            "Phi",
371            type0(),
372            pi(
373                BinderInfo::Default,
374                "Z",
375                type0(),
376                pi(
377                    BinderInfo::Default,
378                    "g",
379                    arrow(bvar(2), arrow(bvar(1), bvar(0))),
380                    pi(
381                        BinderInfo::Default,
382                        "f",
383                        arrow(bvar(1), real_ty()),
384                        arrow(
385                            app(cst("DiffReparameterisation"), bvar(1)),
386                            app2(
387                                cst("Unbiased"),
388                                app2(cst("PathwiseGrad"), bvar(1), bvar(0)),
389                                app2(cst("GradELBO"), app(cst("Reparam"), bvar(1)), bvar(0)),
390                            ),
391                        ),
392                    ),
393                ),
394            ),
395        ),
396    )
397}
398/// **Measure Transport Existence**: for any two probability measures with the same
399/// total mass there exists a measurable transport map.
400///
401/// `measure_transport_exists : ∀ (mu nu : ProbMeasure X),
402///   ∃ (T : X → X), Pushforward T mu = nu`
403pub fn measure_transport_exists_ty() -> Expr {
404    pi(
405        BinderInfo::Default,
406        "X",
407        type0(),
408        pi(
409            BinderInfo::Default,
410            "mu",
411            app(cst("ProbMeasure"), bvar(0)),
412            pi(
413                BinderInfo::Default,
414                "nu",
415                app(cst("ProbMeasure"), bvar(1)),
416                app(
417                    cst("Exists"),
418                    pi(
419                        BinderInfo::Default,
420                        "T",
421                        arrow(bvar(2), bvar(2)),
422                        app2(
423                            cst("Eq"),
424                            app2(cst("Pushforward"), bvar(0), bvar(2)),
425                            bvar(1),
426                        ),
427                    ),
428                ),
429            ),
430        ),
431    )
432}
433/// **Optimal Transport Kantorovich Duality**: the Wasserstein-1 distance equals
434/// the supremum of Lipschitz-1 functions.
435///
436/// `ot_kantorovich : ∀ (mu nu : ProbMeasure X),
437///   W1 mu nu = sup_{f : 1-Lip} (E_mu\[f\] - E_nu\[f\])`
438pub fn ot_kantorovich_ty() -> Expr {
439    pi(
440        BinderInfo::Default,
441        "X",
442        type0(),
443        pi(
444            BinderInfo::Default,
445            "mu",
446            app(cst("ProbMeasure"), bvar(0)),
447            pi(
448                BinderInfo::Default,
449                "nu",
450                app(cst("ProbMeasure"), bvar(1)),
451                app2(
452                    cst("Eq"),
453                    app2(cst("W1Dist"), bvar(1), bvar(0)),
454                    app2(cst("KantorovichDual"), bvar(1), bvar(0)),
455                ),
456            ),
457        ),
458    )
459}
460/// **Stein Identity**: for any smooth function h and score function s_p = ∇ log p,
461/// E_p\[∇ h(x) + h(x) s_p(x)\] = 0.
462///
463/// `stein_identity : ∀ (p : SmoothMeasure X) (h : X → Real),
464///   SmoothTestFn h → E_p\[SteinOp p h\] = 0`
465pub fn stein_identity_ty() -> Expr {
466    pi(
467        BinderInfo::Default,
468        "X",
469        type0(),
470        pi(
471            BinderInfo::Default,
472            "p",
473            app(cst("SmoothMeasure"), bvar(0)),
474            pi(
475                BinderInfo::Default,
476                "h",
477                arrow(bvar(1), real_ty()),
478                arrow(
479                    app(cst("SmoothTestFn"), bvar(0)),
480                    app2(
481                        cst("Eq"),
482                        app2(
483                            cst("Expectation"),
484                            bvar(1),
485                            app2(cst("SteinOp"), bvar(1), bvar(0)),
486                        ),
487                        cst("Real.zero"),
488                    ),
489                ),
490            ),
491        ),
492    )
493}
494/// **Stein Variational Gradient Descent Convergence**: SVGD converges to the target
495/// distribution in the Stein discrepancy sense.
496///
497/// `svgd_convergence : ∀ (target : SmoothMeasure X) (n : Nat),
498///   SteinDiscrepancy (SVGDUpdate target n) target ≤ SVGDBound n`
499pub fn svgd_convergence_ty() -> Expr {
500    pi(
501        BinderInfo::Default,
502        "X",
503        type0(),
504        pi(
505            BinderInfo::Default,
506            "target",
507            app(cst("SmoothMeasure"), bvar(0)),
508            pi(
509                BinderInfo::Default,
510                "n",
511                nat_ty(),
512                app2(
513                    cst("Real.le"),
514                    app2(
515                        cst("SteinDiscrepancy"),
516                        app2(cst("SVGDUpdate"), bvar(1), bvar(0)),
517                        bvar(1),
518                    ),
519                    app(cst("SVGDBound"), bvar(0)),
520                ),
521            ),
522        ),
523    )
524}
525/// **Population Monte Carlo Consistency**: PMC estimators are consistent as
526/// population size and iterations grow.
527///
528/// `pmc_consistency : ∀ (N : Nat) (T : Nat) (target : Measure X),
529///   ConsistentEstimator (PMCEstimator target N T)`
530pub fn pmc_consistency_ty() -> Expr {
531    pi(
532        BinderInfo::Default,
533        "X",
534        type0(),
535        pi(
536            BinderInfo::Default,
537            "N",
538            nat_ty(),
539            pi(
540                BinderInfo::Default,
541                "T",
542                nat_ty(),
543                pi(
544                    BinderInfo::Default,
545                    "target",
546                    app(cst("Measure"), bvar(2)),
547                    app(
548                        cst("ConsistentEstimator"),
549                        app3(cst("PMCEstimator"), bvar(0), bvar(2), bvar(1)),
550                    ),
551                ),
552            ),
553        ),
554    )
555}
556/// **Evolutionary MCMC Detailed Balance**: evolutionary MCMC satisfies detailed
557/// balance with respect to a product-form invariant distribution.
558///
559/// `evol_mcmc_detailed_balance : ∀ (target : Measure X) (temp : Tempering),
560///   DetailedBalance (EvolMCMCKernel target temp) (TemperedTarget target temp)`
561pub fn evol_mcmc_detailed_balance_ty() -> Expr {
562    pi(
563        BinderInfo::Default,
564        "X",
565        type0(),
566        pi(
567            BinderInfo::Default,
568            "target",
569            app(cst("Measure"), bvar(0)),
570            pi(
571                BinderInfo::Default,
572                "temp",
573                cst("Tempering"),
574                app2(
575                    cst("DetailedBalance"),
576                    app2(cst("EvolMCMCKernel"), bvar(1), bvar(0)),
577                    app2(cst("TemperedTarget"), bvar(1), bvar(0)),
578                ),
579            ),
580        ),
581    )
582}
583/// **Parallel Tempering Exchange Correctness**: the swap move in parallel tempering
584/// preserves the joint invariant distribution.
585///
586/// `parallel_tempering_exchange : ∀ (temps : List Real) (joint : Measure (ProductSpace temps)),
587///   InvariantUnder (SwapKernel temps) joint`
588pub fn parallel_tempering_exchange_ty() -> Expr {
589    pi(
590        BinderInfo::Default,
591        "temps",
592        list_ty(real_ty()),
593        pi(
594            BinderInfo::Default,
595            "joint",
596            app(cst("Measure"), app(cst("ProductSpace"), bvar(0))),
597            app2(
598                cst("InvariantUnder"),
599                app(cst("SwapKernel"), bvar(1)),
600                bvar(0),
601            ),
602        ),
603    )
604}
605/// **Simulated Annealing Convergence**: simulated annealing converges to a global
606/// optimum under a logarithmic cooling schedule.
607///
608/// `simulated_annealing_convergence : ∀ (f : X → Real) (T : CoolingSchedule),
609///   LogarithmicSchedule T → ConvergesToGlobalOpt (SAChain f T)`
610pub fn simulated_annealing_convergence_ty() -> Expr {
611    pi(
612        BinderInfo::Default,
613        "X",
614        type0(),
615        pi(
616            BinderInfo::Default,
617            "f",
618            arrow(bvar(0), real_ty()),
619            pi(
620                BinderInfo::Default,
621                "T",
622                cst("CoolingSchedule"),
623                arrow(
624                    app(cst("LogarithmicSchedule"), bvar(0)),
625                    app(
626                        cst("ConvergesToGlobalOpt"),
627                        app2(cst("SAChain"), bvar(2), bvar(0)),
628                    ),
629                ),
630            ),
631        ),
632    )
633}
634/// **VAE ELBO Decomposition**: the VAE objective decomposes as reconstruction term
635/// minus KL divergence.
636///
637/// `vae_elbo_decomp : ∀ (encoder decoder : NeuralNet) (x : X),
638///   VAELBO encoder decoder x = Reconstruction encoder decoder x - KLDivQ encoder x`
639pub fn vae_elbo_decomp_ty() -> Expr {
640    pi(
641        BinderInfo::Default,
642        "X",
643        type0(),
644        pi(
645            BinderInfo::Default,
646            "encoder",
647            cst("NeuralNet"),
648            pi(
649                BinderInfo::Default,
650                "decoder",
651                cst("NeuralNet"),
652                pi(
653                    BinderInfo::Default,
654                    "x",
655                    bvar(2),
656                    app2(
657                        cst("Eq"),
658                        app3(cst("VAELBO"), bvar(2), bvar(1), bvar(0)),
659                        app2(
660                            cst("Real.sub"),
661                            app3(cst("Reconstruction"), bvar(2), bvar(1), bvar(0)),
662                            app2(cst("KLDivQ"), bvar(2), bvar(0)),
663                        ),
664                    ),
665                ),
666            ),
667        ),
668    )
669}
670/// **Diffusion Model Score Matching**: the reverse diffusion score function
671/// minimises the denoising score matching objective.
672///
673/// `diffusion_score_matching : ∀ (p_data : Measure X) (sigma : Real),
674///   OptimScore (DSMObjective p_data sigma) = ScoreFunction (GaussianSmooth p_data sigma)`
675pub fn diffusion_score_matching_ty() -> Expr {
676    pi(
677        BinderInfo::Default,
678        "X",
679        type0(),
680        pi(
681            BinderInfo::Default,
682            "p_data",
683            app(cst("Measure"), bvar(0)),
684            pi(
685                BinderInfo::Default,
686                "sigma",
687                real_ty(),
688                app2(
689                    cst("Eq"),
690                    app(
691                        cst("OptimScore"),
692                        app2(cst("DSMObjective"), bvar(1), bvar(0)),
693                    ),
694                    app(
695                        cst("ScoreFunction"),
696                        app2(cst("GaussianSmooth"), bvar(1), bvar(0)),
697                    ),
698                ),
699            ),
700        ),
701    )
702}
703/// **Flow Matching ODE Correctness**: the conditional flow matching ODE generates
704/// the correct marginal distribution at time t=1.
705///
706/// `flow_matching_ode : ∀ (p_0 p_1 : Measure X) (vt : VectorField),
707///   CondFlowMatchingField vt p_0 p_1 → PushforwardODE vt p_0 1 = p_1`
708pub fn flow_matching_ode_ty() -> Expr {
709    pi(
710        BinderInfo::Default,
711        "X",
712        type0(),
713        pi(
714            BinderInfo::Default,
715            "p_0",
716            app(cst("Measure"), bvar(0)),
717            pi(
718                BinderInfo::Default,
719                "p_1",
720                app(cst("Measure"), bvar(1)),
721                pi(
722                    BinderInfo::Default,
723                    "vt",
724                    cst("VectorField"),
725                    arrow(
726                        app3(cst("CondFlowMatchingField"), bvar(0), bvar(2), bvar(1)),
727                        app2(
728                            cst("Eq"),
729                            app3(cst("PushforwardODE"), bvar(0), bvar(2), cst("Real.one")),
730                            bvar(1),
731                        ),
732                    ),
733                ),
734            ),
735        ),
736    )
737}
738/// **Gaussian Process Posterior**: the posterior of a GP given observations is also a GP.
739///
740/// `gp_posterior_is_gp : ∀ (prior : GaussianProcess X) (obs : Observations),
741///   IsGaussianProcess (GPPosterior prior obs)`
742pub fn gp_posterior_is_gp_ty() -> Expr {
743    pi(
744        BinderInfo::Default,
745        "X",
746        type0(),
747        pi(
748            BinderInfo::Default,
749            "prior",
750            app(cst("GaussianProcess"), bvar(0)),
751            pi(
752                BinderInfo::Default,
753                "obs",
754                cst("Observations"),
755                app(
756                    cst("IsGaussianProcess"),
757                    app2(cst("GPPosterior"), bvar(1), bvar(0)),
758                ),
759            ),
760        ),
761    )
762}
763/// **GP Marginal Likelihood**: the marginal likelihood of a GP is Gaussian.
764///
765/// `gp_marginal_gaussian : ∀ (gp : GaussianProcess X) (X_train : List X),
766///   MarginalLikelihood gp X_train = MultivariateGaussian (GPMean gp X_train) (GPKernelMatrix gp X_train)`
767pub fn gp_marginal_gaussian_ty() -> Expr {
768    pi(
769        BinderInfo::Default,
770        "X",
771        type0(),
772        pi(
773            BinderInfo::Default,
774            "gp",
775            app(cst("GaussianProcess"), bvar(0)),
776            pi(
777                BinderInfo::Default,
778                "X_train",
779                list_ty(bvar(1)),
780                app2(
781                    cst("Eq"),
782                    app2(cst("MarginalLikelihood"), bvar(1), bvar(0)),
783                    app2(
784                        cst("MultivariateGaussian"),
785                        app2(cst("GPMean"), bvar(1), bvar(0)),
786                        app2(cst("GPKernelMatrix"), bvar(1), bvar(0)),
787                    ),
788                ),
789            ),
790        ),
791    )
792}
793/// **Probabilistic Numerics Integration**: Bayesian quadrature produces a posterior
794/// over integrals.
795///
796/// `pn_integration : ∀ (f : X → Real) (prior_gp : GaussianProcess X),
797///   BQPosterior prior_gp f = GaussianMeasureOver Real`
798pub fn pn_integration_ty() -> Expr {
799    pi(
800        BinderInfo::Default,
801        "X",
802        type0(),
803        pi(
804            BinderInfo::Default,
805            "f",
806            arrow(bvar(0), real_ty()),
807            pi(
808                BinderInfo::Default,
809                "prior_gp",
810                app(cst("GaussianProcess"), bvar(1)),
811                app2(
812                    cst("Eq"),
813                    app2(cst("BQPosterior"), bvar(0), bvar(1)),
814                    app(cst("GaussianMeasureOver"), real_ty()),
815                ),
816            ),
817        ),
818    )
819}
820/// **Stein Discrepancy Zero Iff Same Distribution**: the kernel Stein discrepancy
821/// between two measures is zero if and only if they are equal.
822///
823/// `stein_disc_zero_iff : ∀ (p q : Measure X) (k : SteinKernel X),
824///   KSD p q k = 0 ↔ p = q`
825pub fn stein_disc_zero_iff_ty() -> Expr {
826    pi(
827        BinderInfo::Default,
828        "X",
829        type0(),
830        pi(
831            BinderInfo::Default,
832            "p",
833            app(cst("Measure"), bvar(0)),
834            pi(
835                BinderInfo::Default,
836                "q",
837                app(cst("Measure"), bvar(1)),
838                pi(
839                    BinderInfo::Default,
840                    "k",
841                    app(cst("SteinKernel"), bvar(2)),
842                    app2(
843                        cst("Iff"),
844                        app2(
845                            cst("Eq"),
846                            app3(cst("KSD"), bvar(2), bvar(1), bvar(0)),
847                            cst("Real.zero"),
848                        ),
849                        app2(cst("Eq"), bvar(2), bvar(1)),
850                    ),
851                ),
852            ),
853        ),
854    )
855}
856/// **SMC Feynman-Kac**: SMC computes the Feynman-Kac normalising constant exactly
857/// in expectation.
858///
859/// `smc_feynman_kac : ∀ (fk : FeynmanKacModel X) (N : Nat),
860///   E\[SMCNormConst fk N\] = FeynmanKacNormConst fk`
861pub fn smc_feynman_kac_ty() -> Expr {
862    pi(
863        BinderInfo::Default,
864        "X",
865        type0(),
866        pi(
867            BinderInfo::Default,
868            "fk",
869            app(cst("FeynmanKacModel"), bvar(0)),
870            pi(
871                BinderInfo::Default,
872                "N",
873                nat_ty(),
874                app2(
875                    cst("Eq"),
876                    app(
877                        cst("Expectation"),
878                        app2(cst("SMCNormConst"), bvar(1), bvar(0)),
879                    ),
880                    app(cst("FeynmanKacNormConst"), bvar(1)),
881                ),
882            ),
883        ),
884    )
885}
886/// **Particle Marginal Metropolis-Hastings Correctness**: PMMH targeting the exact
887/// posterior is asymptotically exact.
888///
889/// `pmmh_correctness : ∀ (model : LatentModel X Y) (obs : Y),
890///   TargetsExactPosterior (PMMHKernel model obs)`
891pub fn pmmh_correctness_ty() -> Expr {
892    pi(
893        BinderInfo::Default,
894        "X",
895        type0(),
896        pi(
897            BinderInfo::Default,
898            "Y",
899            type0(),
900            pi(
901                BinderInfo::Default,
902                "model",
903                app2(cst("LatentModel"), bvar(1), bvar(0)),
904                pi(
905                    BinderInfo::Default,
906                    "obs",
907                    bvar(1),
908                    app(
909                        cst("TargetsExactPosterior"),
910                        app2(cst("PMMHKernel"), bvar(1), bvar(0)),
911                    ),
912                ),
913            ),
914        ),
915    )
916}
917/// **Annealed Importance Sampling Unbiasedness**: AIS produces unbiased estimates
918/// of the normalising constant.
919///
920/// `ais_unbiased : ∀ (p_0 p_T : Measure X) (beta_sched : AnnealingSchedule),
921///   Unbiased (AISEstimator p_0 p_T beta_sched) (NormConst p_T)`
922pub fn ais_unbiased_ty() -> Expr {
923    pi(
924        BinderInfo::Default,
925        "X",
926        type0(),
927        pi(
928            BinderInfo::Default,
929            "p_0",
930            app(cst("Measure"), bvar(0)),
931            pi(
932                BinderInfo::Default,
933                "p_T",
934                app(cst("Measure"), bvar(1)),
935                pi(
936                    BinderInfo::Default,
937                    "beta_sched",
938                    cst("AnnealingSchedule"),
939                    app2(
940                        cst("Unbiased"),
941                        app3(cst("AISEstimator"), bvar(2), bvar(1), bvar(0)),
942                        app(cst("NormConst"), bvar(1)),
943                    ),
944                ),
945            ),
946        ),
947    )
948}
949/// **Denoising Score Matching Objective**: DSM objective equals implicit score matching
950/// objective under Gaussian noise.
951///
952/// `dsm_equals_sm : ∀ (p_data : Measure X) (sigma : Real),
953///   DSMObjective p_data sigma = ImplicitSMObjective (GaussianConvolution p_data sigma)`
954pub fn dsm_equals_sm_ty() -> Expr {
955    pi(
956        BinderInfo::Default,
957        "X",
958        type0(),
959        pi(
960            BinderInfo::Default,
961            "p_data",
962            app(cst("Measure"), bvar(0)),
963            pi(
964                BinderInfo::Default,
965                "sigma",
966                real_ty(),
967                app2(
968                    cst("Eq"),
969                    app2(cst("DSMObjective"), bvar(1), bvar(0)),
970                    app(
971                        cst("ImplicitSMObjective"),
972                        app2(cst("GaussianConvolution"), bvar(1), bvar(0)),
973                    ),
974                ),
975            ),
976        ),
977    )
978}
979/// **Langevin Dynamics Convergence**: the unadjusted Langevin algorithm (ULA) converges
980/// to the target in 2-Wasserstein under strong convexity.
981///
982/// `langevin_convergence : ∀ (target : LogConcaveMeasure X) (eps : Real),
983///   W2 (ULADist target eps n) target ≤ LangevinBound target eps n`
984pub fn langevin_convergence_ty() -> Expr {
985    pi(
986        BinderInfo::Default,
987        "X",
988        type0(),
989        pi(
990            BinderInfo::Default,
991            "target",
992            app(cst("LogConcaveMeasure"), bvar(0)),
993            pi(
994                BinderInfo::Default,
995                "eps",
996                real_ty(),
997                pi(
998                    BinderInfo::Default,
999                    "n",
1000                    nat_ty(),
1001                    app2(
1002                        cst("Real.le"),
1003                        app3(
1004                            cst("W2"),
1005                            app3(cst("ULADist"), bvar(2), bvar(1), bvar(0)),
1006                            bvar(2),
1007                            bvar(1),
1008                        ),
1009                        app3(cst("LangevinBound"), bvar(2), bvar(1), bvar(0)),
1010                    ),
1011                ),
1012            ),
1013        ),
1014    )
1015}
1016/// **Metropolis-Hastings Detailed Balance**: MH kernel satisfies detailed balance
1017/// w.r.t. the target distribution.
1018///
1019/// `mh_detailed_balance : ∀ (target : Measure X) (proposal : Kernel X X),
1020///   DetailedBalance (MHKernel target proposal) target`
1021pub fn mh_detailed_balance_ty() -> Expr {
1022    pi(
1023        BinderInfo::Default,
1024        "X",
1025        type0(),
1026        pi(
1027            BinderInfo::Default,
1028            "target",
1029            app(cst("Measure"), bvar(0)),
1030            pi(
1031                BinderInfo::Default,
1032                "proposal",
1033                app2(cst("Kernel"), bvar(1), bvar(1)),
1034                app2(
1035                    cst("DetailedBalance"),
1036                    app2(cst("MHKernel"), bvar(1), bvar(0)),
1037                    bvar(1),
1038                ),
1039            ),
1040        ),
1041    )
1042}
1043/// **Gibbs Sampling Invariance**: the Gibbs sampler leaves the joint distribution invariant.
1044///
1045/// `gibbs_invariant : ∀ (joint : Measure (X × Y)),
1046///   InvariantUnder (GibbsKernel joint) joint`
1047pub fn gibbs_invariant_ty() -> Expr {
1048    pi(
1049        BinderInfo::Default,
1050        "X",
1051        type0(),
1052        pi(
1053            BinderInfo::Default,
1054            "Y",
1055            type0(),
1056            pi(
1057                BinderInfo::Default,
1058                "joint",
1059                app(cst("Measure"), app2(cst("Prod"), bvar(1), bvar(0))),
1060                app2(
1061                    cst("InvariantUnder"),
1062                    app(cst("GibbsKernel"), bvar(0)),
1063                    bvar(0),
1064                ),
1065            ),
1066        ),
1067    )
1068}
1069/// **Variational Autoencoder Posterior Collapse**: with a sufficiently expressive decoder
1070/// there exists a risk of posterior collapse.
1071///
1072/// `vae_posterior_collapse_risk : ∀ (decoder : ExpressiveDecoder) (beta : Real),
1073///   beta < 1 → AvoidsCollapse (BetaVAE decoder beta)`
1074pub fn vae_posterior_collapse_risk_ty() -> Expr {
1075    pi(
1076        BinderInfo::Default,
1077        "decoder",
1078        cst("ExpressiveDecoder"),
1079        pi(
1080            BinderInfo::Default,
1081            "beta",
1082            real_ty(),
1083            arrow(
1084                app2(cst("Real.lt"), bvar(0), cst("Real.one")),
1085                app(
1086                    cst("AvoidsCollapse"),
1087                    app2(cst("BetaVAE"), bvar(1), bvar(0)),
1088                ),
1089            ),
1090        ),
1091    )
1092}
1093/// **Gradient of Log Normalizer**: the gradient of the log normalizer of an exponential
1094/// family equals the mean of the sufficient statistics.
1095///
1096/// `grad_log_normalizer : ∀ (eta : NaturalParams) (T : SuffStat),
1097///   Gradient (LogNormalizer T) eta = MeanSuffStat T (ExpFamilyDist T eta)`
1098pub fn grad_log_normalizer_ty() -> Expr {
1099    pi(
1100        BinderInfo::Default,
1101        "T",
1102        type0(),
1103        pi(
1104            BinderInfo::Default,
1105            "eta",
1106            app(cst("NaturalParams"), bvar(0)),
1107            app2(
1108                cst("Eq"),
1109                app(cst("Gradient"), app(cst("LogNormalizer"), bvar(0))),
1110                app2(
1111                    cst("MeanSuffStat"),
1112                    bvar(0),
1113                    app2(cst("ExpFamilyDist"), bvar(0), bvar(1)),
1114                ),
1115            ),
1116        ),
1117    )
1118}
1119/// **Sequential Monte Carlo Genealogy**: the ancestral lineage in SMC traces back
1120/// through the resampling steps.
1121///
1122/// `smc_genealogy : ∀ (pf : ParticleSystem X T),
1123///   AncestralLineage pf = CoalescentProcess (ResamplingTimes pf)`
1124pub fn smc_genealogy_ty() -> Expr {
1125    pi(
1126        BinderInfo::Default,
1127        "X",
1128        type0(),
1129        pi(
1130            BinderInfo::Default,
1131            "T",
1132            nat_ty(),
1133            pi(
1134                BinderInfo::Default,
1135                "pf",
1136                app2(cst("ParticleSystem"), bvar(1), bvar(0)),
1137                app2(
1138                    cst("Eq"),
1139                    app(cst("AncestralLineage"), bvar(0)),
1140                    app(
1141                        cst("CoalescentProcess"),
1142                        app(cst("ResamplingTimes"), bvar(0)),
1143                    ),
1144                ),
1145            ),
1146        ),
1147    )
1148}
1149/// **Kernel Density Estimation Consistency**: the KDE converges to the true density
1150/// in L2 as n → ∞ with optimal bandwidth.
1151///
1152/// `kde_consistency : ∀ (p : SmoothMeasure X) (h : BandwidthSeq),
1153///   OptimalBandwidth h → L2Convergence (KDEn p h)`
1154pub fn kde_consistency_ty() -> Expr {
1155    pi(
1156        BinderInfo::Default,
1157        "X",
1158        type0(),
1159        pi(
1160            BinderInfo::Default,
1161            "p",
1162            app(cst("SmoothMeasure"), bvar(0)),
1163            pi(
1164                BinderInfo::Default,
1165                "h",
1166                cst("BandwidthSeq"),
1167                arrow(
1168                    app(cst("OptimalBandwidth"), bvar(0)),
1169                    app(cst("L2Convergence"), app2(cst("KDEn"), bvar(1), bvar(0))),
1170                ),
1171            ),
1172        ),
1173    )
1174}
1175/// **Variational Inference Mean-Field Factorization**: the mean-field approximation
1176/// optimises each factor holding others fixed via coordinate ascent.
1177///
1178/// `mean_field_cavi : ∀ (joint : Measure Z) (q_factors : List (Measure Z)),
1179///   CAVIStep joint q_factors = UpdatedFactors joint q_factors`
1180pub fn mean_field_cavi_ty() -> Expr {
1181    pi(
1182        BinderInfo::Default,
1183        "Z",
1184        type0(),
1185        pi(
1186            BinderInfo::Default,
1187            "joint",
1188            app(cst("Measure"), bvar(0)),
1189            pi(
1190                BinderInfo::Default,
1191                "q_factors",
1192                list_ty(app(cst("Measure"), bvar(1))),
1193                app2(
1194                    cst("Eq"),
1195                    app2(cst("CAVIStep"), bvar(1), bvar(0)),
1196                    app2(cst("UpdatedFactors"), bvar(1), bvar(0)),
1197                ),
1198            ),
1199        ),
1200    )
1201}
1202/// **Probabilistic Backpropagation Gaussian Propagation**: PBP propagates a
1203/// Gaussian approximation through each layer of a neural network.
1204///
1205/// `pbp_gaussian_propagation : ∀ (net : BayesianNeuralNet) (x : Input),
1206///   GaussianApproxActivations (PBP net x) = PBPActivations net x`
1207pub fn pbp_gaussian_propagation_ty() -> Expr {
1208    pi(
1209        BinderInfo::Default,
1210        "net",
1211        cst("BayesianNeuralNet"),
1212        pi(
1213            BinderInfo::Default,
1214            "x",
1215            cst("Input"),
1216            app2(
1217                cst("Eq"),
1218                app(
1219                    cst("GaussianApproxActivations"),
1220                    app2(cst("PBP"), bvar(1), bvar(0)),
1221                ),
1222                app2(cst("PBPActivations"), bvar(1), bvar(0)),
1223            ),
1224        ),
1225    )
1226}
1227/// **Expectation Propagation Fixed Point**: EP converges when the cavity distribution
1228/// and tilted distribution agree.
1229///
1230/// `ep_fixed_point : ∀ (model : FactorGraph) (approx : GaussianApprox),
1231///   EPFixedPoint approx model ↔ CavityAgreement approx model`
1232pub fn ep_fixed_point_ty() -> Expr {
1233    pi(
1234        BinderInfo::Default,
1235        "model",
1236        cst("FactorGraph"),
1237        pi(
1238            BinderInfo::Default,
1239            "approx",
1240            cst("GaussianApprox"),
1241            app2(
1242                cst("Iff"),
1243                app2(cst("EPFixedPoint"), bvar(0), bvar(1)),
1244                app2(cst("CavityAgreement"), bvar(0), bvar(1)),
1245            ),
1246        ),
1247    )
1248}
1249/// **Nested Monte Carlo Estimator Bias**: nested MC estimators are biased but
1250/// consistent as the inner sample size grows.
1251///
1252/// `nested_mc_bias : ∀ (outer inner : Nat) (f : X → Real),
1253///   Bias (NestedMCEstimator f outer inner) ≤ NestedMCBiasRate inner`
1254pub fn nested_mc_bias_ty() -> Expr {
1255    pi(
1256        BinderInfo::Default,
1257        "X",
1258        type0(),
1259        pi(
1260            BinderInfo::Default,
1261            "outer",
1262            nat_ty(),
1263            pi(
1264                BinderInfo::Default,
1265                "inner",
1266                nat_ty(),
1267                pi(
1268                    BinderInfo::Default,
1269                    "f",
1270                    arrow(bvar(2), real_ty()),
1271                    app2(
1272                        cst("Real.le"),
1273                        app(
1274                            cst("Bias"),
1275                            app3(cst("NestedMCEstimator"), bvar(0), bvar(2), bvar(1)),
1276                        ),
1277                        app(cst("NestedMCBiasRate"), bvar(1)),
1278                    ),
1279                ),
1280            ),
1281        ),
1282    )
1283}
1284/// **Approximate Bayesian Computation Consistency**: ABC-SMC converges to the
1285/// correct posterior as the tolerance ε → 0.
1286///
1287/// `abc_smc_consistency : ∀ (prior : Measure Theta) (sim : Theta → Measure Y) (eps : Real),
1288///   eps > 0 → ApproxPosterior (ABCSMC prior sim eps) eps (TruePosterior prior sim)`
1289pub fn abc_smc_consistency_ty() -> Expr {
1290    pi(
1291        BinderInfo::Default,
1292        "Theta",
1293        type0(),
1294        pi(
1295            BinderInfo::Default,
1296            "Y",
1297            type0(),
1298            pi(
1299                BinderInfo::Default,
1300                "prior",
1301                app(cst("Measure"), bvar(1)),
1302                pi(
1303                    BinderInfo::Default,
1304                    "sim",
1305                    arrow(bvar(2), app(cst("Measure"), bvar(1))),
1306                    pi(
1307                        BinderInfo::Default,
1308                        "eps",
1309                        real_ty(),
1310                        arrow(
1311                            app2(cst("Real.gt"), bvar(0), cst("Real.zero")),
1312                            app3(
1313                                cst("ApproxPosterior"),
1314                                app3(cst("ABCSMC"), bvar(3), bvar(1), bvar(0)),
1315                                bvar(0),
1316                                app2(cst("TruePosterior"), bvar(3), bvar(1)),
1317                            ),
1318                        ),
1319                    ),
1320                ),
1321            ),
1322        ),
1323    )
1324}
1325/// Populate an `Environment` with all probabilistic-programming axiom declarations.
1326pub fn build_probabilistic_programming_env(
1327    env: &mut Environment,
1328) -> Result<(), Box<dyn std::error::Error>> {
1329    let axioms: &[(&str, Expr)] = &[
1330        ("Measure", measure_ty()),
1331        ("SigmaAlgebra", sigma_algebra_ty()),
1332        ("MeasurableSpace", measurable_space_ty()),
1333        ("ProbabilityMonad", probability_monad_ty()),
1334        ("Kernel", kernel_ty()),
1335        ("Sampler", sampler_ty()),
1336        ("PPLProgram", ppl_program_ty()),
1337        ("GradientEstimator", gradient_estimator_ty()),
1338        ("ParticleFilter", particle_filter_ty()),
1339        (
1340            "Joint",
1341            arrow(
1342                app(cst("Measure"), type0()),
1343                arrow(app2(cst("Kernel"), type0(), type0()), type0()),
1344            ),
1345        ),
1346        ("JointMeasure", arrow(type0(), arrow(type0(), type0()))),
1347        (
1348            "Marginal",
1349            arrow(app2(cst("Kernel"), type0(), type0()), type0()),
1350        ),
1351        (
1352            "Posterior",
1353            arrow(
1354                app(cst("Measure"), type0()),
1355                arrow(
1356                    app2(cst("Kernel"), type0(), type0()),
1357                    app(cst("Measure"), type0()),
1358                ),
1359            ),
1360        ),
1361        ("RadonNikodym", arrow(type0(), arrow(type0(), type0()))),
1362        (
1363            "AbsContinuous",
1364            arrow(
1365                app(cst("Measure"), type0()),
1366                arrow(app(cst("Measure"), type0()), prop()),
1367            ),
1368        ),
1369        ("ConsistentEstimator", arrow(type0(), prop())),
1370        (
1371            "ISSampler",
1372            arrow(
1373                arrow(type0(), real_ty()),
1374                arrow(
1375                    app(cst("Measure"), type0()),
1376                    arrow(app(cst("Measure"), type0()), type0()),
1377                ),
1378            ),
1379        ),
1380        (
1381            "ELBO",
1382            arrow(
1383                app(cst("Measure"), type0()),
1384                arrow(type0(), arrow(type0(), real_ty())),
1385            ),
1386        ),
1387        (
1388            "LogMarginalLikelihood",
1389            arrow(type0(), arrow(type0(), real_ty())),
1390        ),
1391        ("MonadLaws", arrow(arrow(type0(), type0()), prop())),
1392        ("Unbiased", arrow(type0(), arrow(type0(), prop()))),
1393        ("ReparamGradient", arrow(arrow(type0(), real_ty()), type0())),
1394        (
1395            "GradExpectation",
1396            arrow(arrow(type0(), real_ty()), arrow(type0(), type0())),
1397        ),
1398        (
1399            "HMCKernel",
1400            arrow(
1401                app(cst("Measure"), type0()),
1402                app2(cst("Kernel"), type0(), type0()),
1403            ),
1404        ),
1405        (
1406            "InvariantUnder",
1407            arrow(
1408                app2(cst("Kernel"), type0(), type0()),
1409                arrow(app(cst("Measure"), type0()), prop()),
1410            ),
1411        ),
1412        ("StateSpaceModel", arrow(type0(), type0())),
1413        ("ConsistentFilteringDist", arrow(type0(), prop())),
1414        ("bayes_measure_theory", bayes_measure_theory_ty()),
1415        ("giry_monad_laws", giry_monad_laws_ty()),
1416        ("is_consistency", is_consistency_ty()),
1417        ("elbo_lower_bound", elbo_lower_bound_ty()),
1418        ("reparam_unbiased", reparam_unbiased_ty()),
1419        ("hmc_invariant", hmc_invariant_ty()),
1420        ("smc_consistency", smc_consistency_ty()),
1421        ("VariationalFamily", arrow(type0(), type0())),
1422        ("LRSchedule", type0()),
1423        ("SVIOptimizer", arrow(type0(), arrow(type0(), type0()))),
1424        ("ConvergesToLocalOptimum", arrow(type0(), prop())),
1425        ("Bijective", arrow(arrow(type0(), type0()), prop())),
1426        ("FlowDensityEq", arrow(type0(), prop())),
1427        ("ParametricMeasure", arrow(type0(), type0())),
1428        ("RegularFamily", arrow(type0(), prop())),
1429        (
1430            "ScoreFnGrad",
1431            arrow(type0(), arrow(arrow(type0(), real_ty()), type0())),
1432        ),
1433        (
1434            "GradELBO",
1435            arrow(type0(), arrow(arrow(type0(), real_ty()), type0())),
1436        ),
1437        (
1438            "DiffReparameterisation",
1439            arrow(arrow(type0(), arrow(type0(), type0())), prop()),
1440        ),
1441        (
1442            "PathwiseGrad",
1443            arrow(
1444                arrow(type0(), arrow(type0(), type0())),
1445                arrow(arrow(type0(), real_ty()), type0()),
1446            ),
1447        ),
1448        (
1449            "Reparam",
1450            arrow(arrow(type0(), arrow(type0(), type0())), type0()),
1451        ),
1452        ("ProbMeasure", arrow(type0(), type0())),
1453        (
1454            "Pushforward",
1455            arrow(
1456                arrow(type0(), type0()),
1457                arrow(app(cst("Measure"), type0()), app(cst("Measure"), type0())),
1458            ),
1459        ),
1460        (
1461            "W1Dist",
1462            arrow(
1463                app(cst("Measure"), type0()),
1464                arrow(app(cst("Measure"), type0()), real_ty()),
1465            ),
1466        ),
1467        (
1468            "KantorovichDual",
1469            arrow(
1470                app(cst("Measure"), type0()),
1471                arrow(app(cst("Measure"), type0()), real_ty()),
1472            ),
1473        ),
1474        ("SmoothMeasure", arrow(type0(), type0())),
1475        ("SmoothTestFn", arrow(arrow(type0(), real_ty()), prop())),
1476        (
1477            "Expectation",
1478            arrow(app(cst("Measure"), type0()), arrow(type0(), real_ty())),
1479        ),
1480        (
1481            "SteinOp",
1482            arrow(
1483                app(cst("SmoothMeasure"), type0()),
1484                arrow(arrow(type0(), real_ty()), type0()),
1485            ),
1486        ),
1487        ("Real.zero", real_ty()),
1488        ("Real.one", real_ty()),
1489        (
1490            "SteinDiscrepancy",
1491            arrow(
1492                app(cst("Measure"), type0()),
1493                arrow(app(cst("Measure"), type0()), real_ty()),
1494            ),
1495        ),
1496        (
1497            "SVGDUpdate",
1498            arrow(
1499                app(cst("SmoothMeasure"), type0()),
1500                arrow(nat_ty(), app(cst("Measure"), type0())),
1501            ),
1502        ),
1503        ("SVGDBound", arrow(nat_ty(), real_ty())),
1504        ("SteinKernel", arrow(type0(), type0())),
1505        (
1506            "KSD",
1507            arrow(
1508                app(cst("Measure"), type0()),
1509                arrow(
1510                    app(cst("Measure"), type0()),
1511                    arrow(app(cst("SteinKernel"), type0()), real_ty()),
1512                ),
1513            ),
1514        ),
1515        ("Iff", arrow(prop(), arrow(prop(), prop()))),
1516        (
1517            "PMCEstimator",
1518            arrow(
1519                app(cst("Measure"), type0()),
1520                arrow(nat_ty(), arrow(nat_ty(), type0())),
1521            ),
1522        ),
1523        ("Tempering", type0()),
1524        (
1525            "EvolMCMCKernel",
1526            arrow(
1527                app(cst("Measure"), type0()),
1528                arrow(cst("Tempering"), app2(cst("Kernel"), type0(), type0())),
1529            ),
1530        ),
1531        (
1532            "TemperedTarget",
1533            arrow(
1534                app(cst("Measure"), type0()),
1535                arrow(cst("Tempering"), app(cst("Measure"), type0())),
1536            ),
1537        ),
1538        (
1539            "DetailedBalance",
1540            arrow(
1541                app2(cst("Kernel"), type0(), type0()),
1542                arrow(app(cst("Measure"), type0()), prop()),
1543            ),
1544        ),
1545        ("ProductSpace", arrow(list_ty(real_ty()), type0())),
1546        (
1547            "SwapKernel",
1548            arrow(list_ty(real_ty()), app2(cst("Kernel"), type0(), type0())),
1549        ),
1550        ("CoolingSchedule", type0()),
1551        ("LogarithmicSchedule", arrow(cst("CoolingSchedule"), prop())),
1552        (
1553            "SAChain",
1554            arrow(
1555                arrow(type0(), real_ty()),
1556                arrow(cst("CoolingSchedule"), type0()),
1557            ),
1558        ),
1559        ("ConvergesToGlobalOpt", arrow(type0(), prop())),
1560        ("NeuralNet", type0()),
1561        (
1562            "VAELBO",
1563            arrow(
1564                cst("NeuralNet"),
1565                arrow(cst("NeuralNet"), arrow(type0(), real_ty())),
1566            ),
1567        ),
1568        (
1569            "Reconstruction",
1570            arrow(
1571                cst("NeuralNet"),
1572                arrow(cst("NeuralNet"), arrow(type0(), real_ty())),
1573            ),
1574        ),
1575        ("KLDivQ", arrow(cst("NeuralNet"), arrow(type0(), real_ty()))),
1576        ("Real.sub", arrow(real_ty(), arrow(real_ty(), real_ty()))),
1577        (
1578            "DSMObjective",
1579            arrow(app(cst("Measure"), type0()), arrow(real_ty(), type0())),
1580        ),
1581        (
1582            "GaussianSmooth",
1583            arrow(
1584                app(cst("Measure"), type0()),
1585                arrow(real_ty(), app(cst("Measure"), type0())),
1586            ),
1587        ),
1588        ("OptimScore", arrow(type0(), type0())),
1589        (
1590            "ScoreFunction",
1591            arrow(app(cst("Measure"), type0()), type0()),
1592        ),
1593        ("VectorField", type0()),
1594        (
1595            "CondFlowMatchingField",
1596            arrow(
1597                cst("VectorField"),
1598                arrow(
1599                    app(cst("Measure"), type0()),
1600                    arrow(app(cst("Measure"), type0()), prop()),
1601                ),
1602            ),
1603        ),
1604        (
1605            "PushforwardODE",
1606            arrow(
1607                cst("VectorField"),
1608                arrow(
1609                    app(cst("Measure"), type0()),
1610                    arrow(real_ty(), app(cst("Measure"), type0())),
1611                ),
1612            ),
1613        ),
1614        ("GaussianProcess", arrow(type0(), type0())),
1615        ("Observations", type0()),
1616        ("IsGaussianProcess", arrow(type0(), prop())),
1617        (
1618            "GPPosterior",
1619            arrow(
1620                app(cst("GaussianProcess"), type0()),
1621                arrow(cst("Observations"), type0()),
1622            ),
1623        ),
1624        (
1625            "MarginalLikelihood",
1626            arrow(
1627                app(cst("GaussianProcess"), type0()),
1628                arrow(list_ty(type0()), app(cst("Measure"), type0())),
1629            ),
1630        ),
1631        (
1632            "MultivariateGaussian",
1633            arrow(type0(), arrow(type0(), app(cst("Measure"), type0()))),
1634        ),
1635        (
1636            "GPMean",
1637            arrow(
1638                app(cst("GaussianProcess"), type0()),
1639                arrow(list_ty(type0()), type0()),
1640            ),
1641        ),
1642        (
1643            "GPKernelMatrix",
1644            arrow(
1645                app(cst("GaussianProcess"), type0()),
1646                arrow(list_ty(type0()), type0()),
1647            ),
1648        ),
1649        (
1650            "BQPosterior",
1651            arrow(
1652                app(cst("GaussianProcess"), type0()),
1653                arrow(arrow(type0(), real_ty()), app(cst("Measure"), real_ty())),
1654            ),
1655        ),
1656        (
1657            "GaussianMeasureOver",
1658            arrow(type0(), app(cst("Measure"), real_ty())),
1659        ),
1660        ("FeynmanKacModel", arrow(type0(), type0())),
1661        (
1662            "SMCNormConst",
1663            arrow(
1664                app(cst("FeynmanKacModel"), type0()),
1665                arrow(nat_ty(), real_ty()),
1666            ),
1667        ),
1668        (
1669            "FeynmanKacNormConst",
1670            arrow(app(cst("FeynmanKacModel"), type0()), real_ty()),
1671        ),
1672        ("LatentModel", arrow(type0(), arrow(type0(), type0()))),
1673        (
1674            "PMMHKernel",
1675            arrow(
1676                app2(cst("LatentModel"), type0(), type0()),
1677                arrow(type0(), app2(cst("Kernel"), type0(), type0())),
1678            ),
1679        ),
1680        (
1681            "TargetsExactPosterior",
1682            arrow(app2(cst("Kernel"), type0(), type0()), prop()),
1683        ),
1684        ("AnnealingSchedule", type0()),
1685        (
1686            "AISEstimator",
1687            arrow(
1688                app(cst("Measure"), type0()),
1689                arrow(
1690                    app(cst("Measure"), type0()),
1691                    arrow(cst("AnnealingSchedule"), type0()),
1692                ),
1693            ),
1694        ),
1695        ("NormConst", arrow(app(cst("Measure"), type0()), real_ty())),
1696        (
1697            "ImplicitSMObjective",
1698            arrow(app(cst("Measure"), type0()), type0()),
1699        ),
1700        (
1701            "GaussianConvolution",
1702            arrow(
1703                app(cst("Measure"), type0()),
1704                arrow(real_ty(), app(cst("Measure"), type0())),
1705            ),
1706        ),
1707        ("LogConcaveMeasure", arrow(type0(), type0())),
1708        (
1709            "ULADist",
1710            arrow(
1711                app(cst("LogConcaveMeasure"), type0()),
1712                arrow(real_ty(), arrow(nat_ty(), app(cst("Measure"), type0()))),
1713            ),
1714        ),
1715        (
1716            "W2",
1717            arrow(
1718                app(cst("Measure"), type0()),
1719                arrow(app(cst("Measure"), type0()), real_ty()),
1720            ),
1721        ),
1722        (
1723            "LangevinBound",
1724            arrow(
1725                app(cst("LogConcaveMeasure"), type0()),
1726                arrow(real_ty(), arrow(nat_ty(), real_ty())),
1727            ),
1728        ),
1729        (
1730            "MHKernel",
1731            arrow(
1732                app(cst("Measure"), type0()),
1733                arrow(
1734                    app2(cst("Kernel"), type0(), type0()),
1735                    app2(cst("Kernel"), type0(), type0()),
1736                ),
1737            ),
1738        ),
1739        ("Prod", arrow(type0(), arrow(type0(), type0()))),
1740        (
1741            "GibbsKernel",
1742            arrow(
1743                app(cst("Measure"), app2(cst("Prod"), type0(), type0())),
1744                app2(
1745                    cst("Kernel"),
1746                    app2(cst("Prod"), type0(), type0()),
1747                    app2(cst("Prod"), type0(), type0()),
1748                ),
1749            ),
1750        ),
1751        ("ExpressiveDecoder", type0()),
1752        (
1753            "BetaVAE",
1754            arrow(cst("ExpressiveDecoder"), arrow(real_ty(), type0())),
1755        ),
1756        ("AvoidsCollapse", arrow(type0(), prop())),
1757        ("Real.lt", arrow(real_ty(), arrow(real_ty(), prop()))),
1758        ("NaturalParams", arrow(type0(), type0())),
1759        ("LogNormalizer", arrow(type0(), arrow(type0(), real_ty()))),
1760        ("Gradient", arrow(arrow(type0(), real_ty()), type0())),
1761        (
1762            "MeanSuffStat",
1763            arrow(type0(), arrow(app(cst("Measure"), type0()), type0())),
1764        ),
1765        (
1766            "ExpFamilyDist",
1767            arrow(type0(), arrow(type0(), app(cst("Measure"), type0()))),
1768        ),
1769        ("ParticleSystem", arrow(type0(), arrow(nat_ty(), type0()))),
1770        ("AncestralLineage", arrow(type0(), type0())),
1771        ("CoalescentProcess", arrow(type0(), type0())),
1772        ("ResamplingTimes", arrow(type0(), type0())),
1773        ("BandwidthSeq", type0()),
1774        (
1775            "KDEn",
1776            arrow(
1777                app(cst("SmoothMeasure"), type0()),
1778                arrow(cst("BandwidthSeq"), type0()),
1779            ),
1780        ),
1781        ("OptimalBandwidth", arrow(cst("BandwidthSeq"), prop())),
1782        ("L2Convergence", arrow(type0(), prop())),
1783        (
1784            "CAVIStep",
1785            arrow(
1786                app(cst("Measure"), type0()),
1787                arrow(
1788                    list_ty(app(cst("Measure"), type0())),
1789                    list_ty(app(cst("Measure"), type0())),
1790                ),
1791            ),
1792        ),
1793        (
1794            "UpdatedFactors",
1795            arrow(
1796                app(cst("Measure"), type0()),
1797                arrow(
1798                    list_ty(app(cst("Measure"), type0())),
1799                    list_ty(app(cst("Measure"), type0())),
1800                ),
1801            ),
1802        ),
1803        ("BayesianNeuralNet", type0()),
1804        ("Input", type0()),
1805        (
1806            "PBP",
1807            arrow(cst("BayesianNeuralNet"), arrow(cst("Input"), type0())),
1808        ),
1809        ("GaussianApproxActivations", arrow(type0(), type0())),
1810        (
1811            "PBPActivations",
1812            arrow(cst("BayesianNeuralNet"), arrow(cst("Input"), type0())),
1813        ),
1814        ("FactorGraph", type0()),
1815        ("GaussianApprox", type0()),
1816        (
1817            "EPFixedPoint",
1818            arrow(cst("GaussianApprox"), arrow(cst("FactorGraph"), prop())),
1819        ),
1820        (
1821            "CavityAgreement",
1822            arrow(cst("GaussianApprox"), arrow(cst("FactorGraph"), prop())),
1823        ),
1824        (
1825            "NestedMCEstimator",
1826            arrow(
1827                arrow(type0(), real_ty()),
1828                arrow(nat_ty(), arrow(nat_ty(), type0())),
1829            ),
1830        ),
1831        ("Bias", arrow(type0(), real_ty())),
1832        ("NestedMCBiasRate", arrow(nat_ty(), real_ty())),
1833        (
1834            "ABCSMC",
1835            arrow(
1836                app(cst("Measure"), type0()),
1837                arrow(
1838                    arrow(type0(), app(cst("Measure"), type0())),
1839                    arrow(real_ty(), app(cst("Measure"), type0())),
1840                ),
1841            ),
1842        ),
1843        (
1844            "TruePosterior",
1845            arrow(
1846                app(cst("Measure"), type0()),
1847                arrow(
1848                    arrow(type0(), app(cst("Measure"), type0())),
1849                    app(cst("Measure"), type0()),
1850                ),
1851            ),
1852        ),
1853        (
1854            "ApproxPosterior",
1855            arrow(
1856                app(cst("Measure"), type0()),
1857                arrow(real_ty(), arrow(app(cst("Measure"), type0()), prop())),
1858            ),
1859        ),
1860        ("Real.gt", arrow(real_ty(), arrow(real_ty(), prop()))),
1861        ("svi_convergence", svi_convergence_ty()),
1862        ("normalizing_flow_cov", normalizing_flow_cov_ty()),
1863        ("score_fn_unbiased", score_fn_unbiased_ty()),
1864        (
1865            "pathwise_gradient_unbiased",
1866            pathwise_gradient_unbiased_ty(),
1867        ),
1868        ("measure_transport_exists", measure_transport_exists_ty()),
1869        ("ot_kantorovich", ot_kantorovich_ty()),
1870        ("stein_identity", stein_identity_ty()),
1871        ("svgd_convergence", svgd_convergence_ty()),
1872        ("pmc_consistency", pmc_consistency_ty()),
1873        (
1874            "evol_mcmc_detailed_balance",
1875            evol_mcmc_detailed_balance_ty(),
1876        ),
1877        (
1878            "parallel_tempering_exchange",
1879            parallel_tempering_exchange_ty(),
1880        ),
1881        (
1882            "simulated_annealing_convergence",
1883            simulated_annealing_convergence_ty(),
1884        ),
1885        ("vae_elbo_decomp", vae_elbo_decomp_ty()),
1886        ("diffusion_score_matching", diffusion_score_matching_ty()),
1887        ("flow_matching_ode", flow_matching_ode_ty()),
1888        ("gp_posterior_is_gp", gp_posterior_is_gp_ty()),
1889        ("gp_marginal_gaussian", gp_marginal_gaussian_ty()),
1890        ("pn_integration", pn_integration_ty()),
1891        ("stein_disc_zero_iff", stein_disc_zero_iff_ty()),
1892        ("smc_feynman_kac", smc_feynman_kac_ty()),
1893        ("pmmh_correctness", pmmh_correctness_ty()),
1894        ("ais_unbiased", ais_unbiased_ty()),
1895        ("dsm_equals_sm", dsm_equals_sm_ty()),
1896        ("langevin_convergence", langevin_convergence_ty()),
1897        ("mh_detailed_balance", mh_detailed_balance_ty()),
1898        ("gibbs_invariant", gibbs_invariant_ty()),
1899        (
1900            "vae_posterior_collapse_risk",
1901            vae_posterior_collapse_risk_ty(),
1902        ),
1903        ("grad_log_normalizer", grad_log_normalizer_ty()),
1904        ("smc_genealogy", smc_genealogy_ty()),
1905        ("kde_consistency", kde_consistency_ty()),
1906        ("mean_field_cavi", mean_field_cavi_ty()),
1907        ("pbp_gaussian_propagation", pbp_gaussian_propagation_ty()),
1908        ("ep_fixed_point", ep_fixed_point_ty()),
1909        ("nested_mc_bias", nested_mc_bias_ty()),
1910        ("abc_smc_consistency", abc_smc_consistency_ty()),
1911    ];
1912    for (name, ty) in axioms {
1913        env.add(Declaration::Axiom {
1914            name: Name::str(*name),
1915            univ_params: vec![],
1916            ty: ty.clone(),
1917        })
1918        .ok();
1919    }
1920    Ok(())
1921}
oxilean_std/probabilistic_programming/functions.rs

oxilean_std/probabilistic_programming/
functions.rs
