Guillaume Baudart

Javier Burroni

Martin Hirzel

Louis Mandel

Avraham Shinnar

Compiling Stan to generative probabilistic languages and extension to deep probabilistic programming

Keywords: CCS Concepts:, Software and its engineering → Compilers;, Theory of computation → Probabilistic computation Probabilistic programming, Semantics, Stan, Pyro ,[e n 1 ,...,e n m ]] | e

published or not. The documents may come   L'archive ouverte pluridisciplinaire

Introduction

Probabilistic Programming Languages (PPLs) are designed to describe probabilistic models and run inference on them. There exists a variety of PPLs. BUGS [START_REF] Lunn | The BUGS project: Evolution, critique and future directions[END_REF], JAGS [START_REF] Plummer | JAGS: A program for analysis of Bayesian graphical models using Gibbs sampling[END_REF], and Stan [START_REF] Carpenter | Stan: A probabilistic programming language[END_REF] focus on efficiency, constraining what is expressible to a subset of models which support fast inference techniques. These languages enjoy broad adoption by the statistics and social sciences communities [START_REF] Bradley | Bayesian methods for data analysis[END_REF][START_REF] Gelman | Data analysis using regression and multilevel/hierarchical models[END_REF][START_REF] Gelman | Bayesian data analysis[END_REF]. Generative languages like Church [START_REF] Goodman | Church: a language for generative models[END_REF], Anglican [START_REF] Tolpin | Design and Implementation of Probabilistic Programming Language Anglican[END_REF], WebPPL [START_REF] Goodman | The Design and Implementation of Probabilistic Programming Languages[END_REF], Pyro [START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF], and Gen [START_REF] Marco | Gen: a general-purpose probabilistic programming system with programmable inference[END_REF] describe generative models, i.e., stochastic procedures that simulate the data generation process. Coming from a core programming languages heritage, generative PPLs typically support rich control constructs and models over structured data. Generative PPLs are increasingly used in machinelearning research and are rapidly incorporating new ideas, such as Stochastic Gradient Variational Inference (SVI), in what is now called Deep Probabilistic Programming [START_REF] Baudart | Deep Probabilistic Programming Languages: A Qualitative Study[END_REF][START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF][START_REF] Tran | Deep Probabilistic Programming[END_REF].

While the semantics of probabilistic languages have been extensively studied [START_REF] Gordon | Probabilistic programming[END_REF][START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF][START_REF] Kozen | Semantics of Probabilistic Programs[END_REF][START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF], to the best of our knowledge little is known about the relationship between Stan and generative PPLs. We show that a simple 1:1 translation is incorrect or incomplete for a set of subtle but widely-used Stan features, such as left expressions or implicit priors.

This paper formalizes the relationship between Stan and generative PPLs and introduces, with correctness proof, a comprehensive compilation scheme that can compile any Stan program to a generative PPL. This enables leveraging the rich set of existing Stan models for testing, benchmarking, or experimenting with new features or inference techniques. Based on this compilation scheme we implemented two new backends for the Stanc3 compiler targeting Pyro [START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF] and NumPyro [START_REF] Phan | Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro[END_REF], a JAX [START_REF] Bradbury | JAX: composable transformations of Python+NumPy programs[END_REF] based version of Pyro. Both Pyro and NumPyro runtimes offer NUTS [START_REF] Hoffman | The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo[END_REF] (No U-Turn Sampler), an optimized Hamiltonian Monte-Carlo (HMC) algorithm that is the preferred inference method for Stan. We can thus validate our approach against Stan. Results show that models compiled using our NumPyro backend yield equivalent results while being 2.3x faster than their Stan counterpart in the geometric mean over 26 benchmarks. Our compiler and runtime library are open-source at https://github.com/deepppl.

In addition, recent probabilistic languages offer new features to program and reason about complex models. Our compilation scheme combined with conservative extensions of Stan can be used to make these benefits available to Stan users. As a proof of concept, based on our Pyro backend, this paper introduces DeepStan: Stan extended with support for explicit variational guides and deep probabilistic models. Variational inference was central in the design of Pyro and programmers can easily craft their own inference guides to run variational inference on probabilistic models. Pyro is built on top of PyTorch [START_REF] Paszke | Automatic Differentiation in PyTorch[END_REF]. Programmers can thus seamlessly import neural networks designed with the state-of-the-art API provided by PyTorch. This paper makes the following contributions:

• A comprehensive compilation scheme from Stan to a generative PPL (Section 2). • Correctness proof of the compilation scheme (Section 3).

• An open-source implementation of two new backends for Stanc3 targeting Pyro and NumPyro (Section 4). • An extension of Stan with explicit variational inference guides and deep probabilistic models (Section 5).

The fundamental new result of this paper is that every Stan program can be expressed as a generative probabilistic program. Besides advancing the understanding of probabilistic programming languages at a fundamental level, this paper aims to provide practical benefits to the communities of both Stan and Pyro. From the perspective of the Stan community, this paper presents a new competitive compiler backend and additional capabilities while retaining familiar syntax and semantics. This compiler can thus be used to migrate existing Stan codebases to Pyro and NumPyro. From the perspective of the Pyro community, this paper presents a new compiler frontend that unlocks many existing real-world models as examples and benchmarks. An extended version of the paper with appendices presenting the proofs and the evaluation results is available [START_REF] Baudart | Compiling Stan to Generative Probabilistic Languages and Extension to Deep Probabilistic Programming[END_REF].

Overview

This section shows how to compile Stan [START_REF] Carpenter | Stan: A probabilistic programming language[END_REF], which specifies a joint probability distribution, to a generative PPL like Church, Anglican, or Pyro. This translation also demonstrates that Stan's expressive power is at most as large as that of generative languages, a fact that was not clear before our paper.

As a running example, consider the biased coin model in Figure 1. Stan's data block defines observed variables for N coin flips x i , i ∈ [1 : N ], which can be 0 for tails or 1 for heads. The parameters block introduces a latent variable z ∈ [0, 1] for the bias of the coin. The model block sets the prior of the bias z to Beta(1, 1), i.e., a uniform distribution over [0,[START_REF] Baudart | Compiling Stan to Generative Probabilistic Languages and Extension to Deep Probabilistic Programming[END_REF]. The for loop indicates that coin flips x i are independent and identically distributed (IID) and depend on z via a Bernoulli distribution. Given concrete observed coin flips, inference yields a posterior distribution for z conditioned on x 1 , . . . , x N .

Generative Translation

Generative PPLs are general-purpose languages extended with two probabilistic constructs [START_REF] Gordon | Probabilistic programming[END_REF][START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF][START_REF] Van De Meent | An Introduction to Probabilistic Programming[END_REF]: sample(D) generates a sample from a distribution D and factor(v) assigns a score v to the current execution trace. Typically, factor is used to condition the model on input data [START_REF] Tolpin | Design and Implementation of Probabilistic Programming Language Anglican[END_REF]. We also introduce observe(D,v) as a syntactic shortcut for factor(D pdf (v)) where D pdf denotes the probability density function of D. This construct penalizes executions according to the score of v w.r.t. D which captures the assumption that the observed data v follows the distribution D.

Compilation. Stan uses the same syntax v ~D for both observed and latent variables. The distinction comes from the kind of the left-hand-side variable: observed variables are declared in the data block, latent variables are declared in the parameters block. A straightforward generative translation compiles a statement v ~D into v = sample(D) if v is a parameter or observe(D, v) if v is data. For example, Figure 2a shows the compiled (using the generative scheme) version of the Stan model of Figure 1 in Python syntax.

Non-generative Features

In Stan, a model represents the unnormalized density of the joint distribution of the parameters defined in the parameters block given the data defined in the data block [START_REF] Carpenter | Stan: A probabilistic programming language[END_REF][START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF]. A Stan program can thus be viewed as a function from parameters Table 1. Stan features that defy generative translation: prevalence, example, and compilation.

Feature

% Example Compilation

Left expression 15 sum(phi) ~normal(0, 0.001*N); observe(Normal(0.,0.001*N), sum(phi))

Multiple updates 8

phi_y ~normal(0,sigma_py); phi_y ~normal(0,sigma_pt) observe(Normal(0.,sigma_py), phi_y); observe(Normal(0.,sigma_pt), phi_y)

Implicit prior 58

real alpha0; /* missing 'alpha0 ~...' */ alpha0 = sample(improper_uniform())

and data to the value of a special variable target that represents the log-density of the model. A Stan model can be described using classic imperative statements, plus two special statements that modify the value of target. The first one, target+= e, increments the value of target by e. The second one, e ~D, is equivalent to target+= D lpdf (e) [START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF] where D lpdf denotes the log probability density function of D.

Unfortunately, these constructs allow the definition of models that cannot be translated using the generative translation defined above. Table 1 lists the Stan features that are not handled correctly. A left expression is where the left-handside of ~is an arbitrary expression. Multiple updates occur when the same parameter appears on the left-hand-side of multiple ~statements. An implicit prior occurs when there is no explicit ~statement in the model for a parameter.

The ł%ž column of Table 1 indicates the percentage of Stan models that exercise each of the non-generative features among the 531 valid files in https://github.com/standev/example-models. The example column contains illustrative excerpts from such models. Since these are official and long-standing examples, we assume that they use the non-generative features on purpose. Comments in the source code further corroborate that the programmer knowingly used the features. While some features only occur in a minority of models, their prevalence is too high to ignore.

Comprehensive Translation

The previous section illustrates that Stan is centered around the definition of target, not around generating samples for parameters, which is required by generative PPLs. The comprehensive translation adds an initialization step to generate samples for all the parameters and compiles all Stan ~statements as observations. Parameter initialization draws from the uniform distribution in their definition domain. For the biased coin example, the result of this translation is shown in Figure 2b: The parameter z is first sampled uniformly on its definition domain and then conditioned with an observation.

The compilation column of Table 1 illustrates the translation of non-generative features. Left expression and multiple updates are simply compiled into observations. Parameter initialization uses the uniform distribution over its definition domain. For unbounded domains, we introduce new distributions (e.g., improper_uniform) with a constant density that can be normalized away during inference. Section 3.3 details the complete compilation scheme.

Intuition of correctness. The semantics of Stan as described in [START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF] is a classic imperative semantics. Its environment includes the special variable target, the unnormalized log-density of the model. On the other hand, the semantics of a generative PPL as described in [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF] defines a kernel mapping an environment to a measurable function. Our compilation scheme adds uniform initializations for all parameters which comes down to the Lebesgue measure on the parameters space, and translates all ~statements to observe statements. We can then show that a succession of observe statements yields a distribution with the same log-density as the Stan semantics. Section 3.4 details the correctness proof.

Implementation. The comprehensive compilation scheme can compile any Stan program to a generative PPL. Section 4 discusses the implementation of two new backends for the Stanc3 compiler targeting Pyro [START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF] ś a PPL in the line of WebPPL [START_REF] Goodman | The Design and Implementation of Probabilistic Programming Languages[END_REF] ś and NumPyro ś a JAX [START_REF] Bradbury | JAX: composable transformations of Python+NumPy programs[END_REF] based version of Pyro. Section 6 experimentally validates that our backends can compile most existing Stan models. Results also show that models compiled using our NumPyro backend outperform their Stan counterpart on existing benchmarks.

Extensions. Pyro is a deep universal probabilistic programming languages with native support for variational inference. Building on Pyro, we use our compiler to extend Stan with support for explicit variational guides (Section 5.1) and deep neural networks to capture complex relations between parameters (Section 5.2).

Semantics and Compilation

This section, building on previous work, first formally defines the semantics of Stan (Section 3.1) and the semantics of GProb, a small generative probabilistic language (Section 3.2). It then defines the compilation function from Stan to GProb (Section 3.3) and proves its correctness (Section 3.4).

Stan: A Declarative Probabilistic Language

The Stan language is informally described in [START_REF] Carpenter | Stan: A probabilistic programming language[END_REF]. A Stan program is a sequence of blocks which in order: declares functions, declares input names and types, pre-processes input data, declares the parameters to infer, defines transformations on parameters, defines the model, and post-processes the parameters. The only mandatory block is model. Variables declared in a block are visible in subsequent blocks. Variable declarations (decl * ) are lists of variables names with their types (e.g., int N;) or arrays with their type and shape (e.g., real x[N]). Types can be constrained to specify the domain of a variable (e.g., real <lower=0> x for x ∈ R + ). Note that vector and matrix are primitive types that can be used in arrays (e.g. vector [N] x [START_REF] Gelman | Data analysis using regression and multilevel/hierarchical models[END_REF] is an array of 10 vectors of size N). Shapes and sizes of arrays, matrices, and vectors are explicit and can be arbitrary expressions. Semantics. Stan programs are evaluated in three steps:

1. data preparation with data and transformed data 2. inference over the model defined by parameters, transformed parameters, and model 3. post-processing with generated quantities.

Sections transformed data, transformed parameters, and generated quantities introduce new variables. Semantically, these sections can all be inlined in the model section. Any Stan program can thus be rewritten into an equivalent program with only the three blocks data, parameters, and model. Alternatively, we show in Section 3.3 that the transformed data section can be pre-computed and passed as input to the model, and the generated quantities can be post-processed after the inference. Functions declared in functions are inlined (stmt ′ is equivalent to stmt after inlining). To simplify the presentation, we focus on this simplified language.

Notations. To refer to the different parts of a program, we will use the following functions. In the following, an environment γ : Var → Val is a mapping from variables to values, γ (x) returns the value of the variable x in an environment γ , γ [x ← v] returns the environment γ where the value of x is set to v, and γ 1 , γ 2 denotes the union of two environments.

The notation ∫ X µ(dx)f (x) is the integral of f w.r.t. the measure µ where x ∈ X is the integration variable. When µ is the Lebesgue measure we also write ∫ X f (x)dx. Following [START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF], we define the semantics of the model block as a deterministic function that takes an initial environment containing the input data and the parameters, and returns an updated environment where the value of the variable target is the un-normalized log-density of the model.

We can then define the semantics of a Stan program as a kernel [START_REF] Kozen | Semantics of Probabilistic Programs[END_REF][START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF][START_REF] Staton | Semantics for probabilistic programming: higher-order functions, continuous distributions, and soft constraints[END_REF], that is, a function {[p]} : D → Σ X → [0, ∞] where Σ X denotes the σ -algebra of the parameter domain X , that is, the set of measurable sets of the product space of parameter values. Given an environment D containing the input data, p D is a measure that maps a measurable set of parameter values U to a score in [0, ∞] obtained by integrating the density of the model, exp(target), over all the possible parameter values in U .

{[p]} D = λU . ∫ U exp( model(p) D[params(p)←θ ] (target)) dθ x = e γ = γ [x ← e γ ]
x We now detail the semantics of statements and expressions in a model block. This formalization is similar to the semantics proposed in [START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF] but expressed denotationally.

Statements. The semantics of a statement s : (Var → Val) → (Var → Val) is a function from an environment γ to an updated environment. Figure 3 gives the semantics of Stan statements. The initial environment contains the input data, the parameters, and the reserved variable target initialized to 0. An assignment updates the value of a variable or of a cell of an indexed structure in the environment. A sequence s 1 ; s 2 evaluates s 2 in the environment produced by s 1 . A for loop on ranges first evaluates the value of the bounds n 1 and n 2 and then repeats the execution of the body 1 + n 2n 1 times. Iterations over indexed structures (for (x in e) {s}) are syntactic sugar over loops on ranges. The behavior depends on the underlying type. For vectors and arrays, iteration is limited to one dimension.

for (x in e) {s} γ = let v = e γ in (i is a fresh variable) for (i in 1:length(v)) {x = v[i]; s} γ
For matrices, iteration is over the two dimensions:

for (x in e) {s} γ = (i and j are fresh variables)

let v = e γ in for (i in 1:length(v)) for (j in 1:length(v[i])) {x = v[i][j]; s} γ
A while loop repeats the execution of its body while the condition is not 0. Expressions. The semantics of an expression e : (Var → Val) → Val is a function from a environment to values. Figure 4 gives the semantics of Stan expressions. Constants evaluate to themselves. Variables are looked up in the environment. Arrays, vectors, and matrix expressions evaluate all their components. Indexing expressions obtain the corresponding value in the associated data. Function calls apply the function to the value of the arguments. Functions are built-ins like + or normal (user-defined functions are inlined).

Limitations. We consider only terminating programs which means in particular that all loops perform a bounded number of iterations. We also limit the access and update of target to the statements target += e and e 1 ~e2 . Assumption 1. All programs terminate.

Assumption 2. Expressions cannot depend on target.

GProb: A Simple Generative PPL

To formalize the compilation, we first define the target language: GProb, a simple generative probabilistic language similar to the one defined in [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF]. GProb is an expression language with the following syntax: An expression is either a Stan expression, a local binding (let), a conditional (if), or a loop (for or while). To simplify the presentation, loops are parameterized by the set X of variables updated and returned by their body. Moreover, we limit the definition of the semantics to terminating loops that do not depend on sampled values in the body of the loop. GProb also contains the classic probabilistic expressions: sample draws a sample from a distribution, and factor assigns a score to the current execution trace to condition the model. The return expression lifts a deterministic expression to a probabilistic context. We also introduce observe(D,v) as a syntactic shortcut for factor(D pdf (v)) where D pdf denotes the density function of D. This construct penalizes the current execution with the likelihood of v w.r.t. D which captures the assumption that the observed data v follows the distribution D.

{[return(e)]} γ = λU . δ e γ (U ) {[let x = e 1 in e 2 ]} γ = λU . ∫ X {[e 1 ]} γ (dv) × {[e 2 ]} γ [x ←v] (U ) {[let x[e 1 ,...,e n ] = e in e ′ ]} γ = λU . ∫ X {[e]} γ (dv) × {[e ′ ]} γ [x ←(x [ e 1 γ ,..., e n γ ]←v)] (U ) {[for X (x in e 1 :e 2 ) e 3 ]} γ = λU .let n 1 = e 1 γ in let n 2 = e 2 γ in if n 1 > n 2 then δ γ (X) (U ) else ∫ X {[e 3 ]} γ [x ←n 1 ] (dv) × {[for X (x in n 1 + 1:n 2 ) e 3 ]} γ [X←v] (U ) {[while X (e 1 ) e 2 ]} γ = λU .if e 1 γ = 0 then δ γ (X) (U ) else ∫ X {[e 2 ]} γ (dv) × {[while X (e 1 ) e 2 ]} γ [X←v] (U ) {[if (e) e 1 else e 2 ]} γ = λU .if e γ 0 then {[e 1 ]} γ (U ) else {[e 2 ]} γ (U ) {[sample(e)]} γ = λU . e γ (U ) {[factor(e)]} γ = λU . exp( e γ )δ () (U )

Figure 5. Generative probabilistic language semantics

Semantics. Following [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF] we give a measure-based semantics to GProb. The semantics of an expression is a kernel that given an environment returns a measure on the set of possible values. Given input data, normalizing the corresponding measure computes the program's posterior distribution.

The semantics of GProb is given in Figure 5. A deterministic expression is lifted to a probabilistic expression with the Dirac delta measure (δ x (U ) = 1 if x ∈ U , 0 otherwise). A local definition let x = e 1 in e 2 is interpreted by integrating the semantics of e 2 over the set of all possible values for x.

Compared to the language defined in [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF], we added Stanlike loops. Loops behave like a sequence of expressions and return the values of the variables updated by their body. We impose that the condition of a loop cannot depend on parameters sampled in the loop body, and consider only terminating loops. Hence for any given context γ , it suffices to unroll the definition of the loop semantics a finite number of times to get a measure term describing the semantics.

Finally, the semantics of probabilistic operators is the following. The semantics of sample(e) is the probability distribution e γ (e.g. N (0, 1)). A type system, omitted here for conciseness, ensures that we only sample from distributions. The semantics of factor(e) is a measure defined on

C k (t cstr x;) = let x = sample(C (cstr, [])) in k C k (t cstr x[shape];) = let x = sample(C (cstr, shape)) in k C k (decl decls) = C C k (decls) (decl) C (ε, shape) = improper_uniform([-∞, ∞],shape) C (<lower=e 1 >, shape) = improper_uniform([e 1 , ∞],shape) C (<upper=e 2 >, shape) = improper_uniform([-∞, e 2 ],shape) C (<lower=e 1 ,upper=e 2 >, shape) = uniform([e 1 , e 2 ],shape) Figure 6. Comprehensive compilation of parameters C k (x = e) = let x = return(e) in k C k (x[e 1 ,...,e n ] = e) = let x[e 1 ,...,e n ] = e in k C k (s 1 ; s 2 ) = C C k (s 2 ) (s 1 ) C k (for (x in e 1 :e 2 ) {s}) = let lhs(s) = for lhs(s) (x in e 1 :e 2 ) C return(lhs(s)) (s) in k C k (while (e) {s}) = let lhs(s) = while lhs(s) (e) C return(lhs(s)) (s) in k C k (if (e) s 1 else s 2 ) = if (e) C k (s 1 ) else C k (s 2 ) C k (skip) = k C k (target += e) = let () = factor(e) in k C k (e ~f (e 1 ,...,e n )) = let () = observe(f (e 1 ,...,e n ),e) in k Figure 7.
Comprehensive compilation of statements the singleton space () whose value is exp( e ) (this operator corresponds to score in [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF] but in log-scale, which is common for numerical precision).

Comprehensive Translation

The key idea is to first sample all parameters from priors with a constant density that can be normalized away during inference (e.g., uniform on bounded domains), and then compile all ~statements into observe statements. The compilation functions C k (params(p)) for the parameters and C k (model(p)) for the model are both parameterized by a continuation k. The compilation of the entire program first compiles the parameters to introduce the priors, then compiles the model, and finally adds a return statement for all the parameters. In continuation passing style:

C(p) = C C return(params(p)) (model(p)) (params(p))
Parameters. In Stan, parameters are defined on R n with optional domain constraints (e.g. <lower=0>). For each parameter, the comprehensive translation sets the prior to either the uniform distribution on a bounded domain, or an improper prior with a constant density w.r.t. the Lebesgue measure that we call improper_uniform. The compilation function of the parameters, defined Figure 6, thus produces a succession of sample expressions:

C k (params(p)) = let x 1 = D 1 in . . . let x n = D n in k
where each D i is either uniform or improper_uniform. The compilation of Stan loops raises an issue. In Stan, all the variables that appear on the left-hand side of an assignment in the body of a loop are state variables that are incrementally updated at each iteration. Since GProb is a functional language, the state of the loops is made explicit. To propagate the environment updates at each iteration, loop expressions are annotated with all the variables that are assigned in their body (lhs(stmt)). These variables are returned at the end of the loop and can be used in the continuation.

Statements.

Pre-and post-processing blocks. Section 3.1 showed that all Stan programs can be rewritten in a kernel using only the data, parameters, and model sections. This approach can make the model more complicated and thus, the inference more expensive. In particular, pre-and post-processing steps do not need to be computed at each inference step.

In Stan, users can define functions in the functions block. These functions can be compiled into functions in the target language using the comprehensive translation.

The code in the transformed data section only depends on variables declared in the data section and can be computed only once before the inference. We compile this section into a function that takes as argument the data and returns the transformed data. The variables declared in the transformed data section become new inputs for the model.

On the other hand, the transformed parameters block must be part of the model since it depends on the model parameters. This section is thus inlined in the compiled model.

Finally, the generated quantities block is executed only once on the inference result. It is compiled into a function of the data, transformed data, and parameters returned by the model. The transformed parameters block is also inlined, since generated quantities may depend on them.

Correctness of the Compilation

We can now show that a Stan program and the corresponding compiled code yield the same un-normalized measure up to a constant factor (and thus the same posterior distribution). The proof has two steps: (1) simplifying the sequence of sample statements introduced by the compilation of the parameters, and [START_REF] Baudart | Deep Probabilistic Programming Languages: A Qualitative Study[END_REF] showing that the value of the Stan target corresponds to the score computed by the generated code.

Priors. First, we simplify the nested integrals introduced by the sequence of sample statements for the parameters priors into one integral over the parameter domain.

Lemma 3.1. For all Stan programs p with stmt = model(p) and P = params(p), and environments γ :

{[C(p)]} γ ∝ λU . ∫ U {[C return(()) (stmt)]} γ [P←θ ] ({()})dθ
where U ∈ Σ X is a measurable set of parameter values, with X = Dom(P).

The proof relies on the fact that the parameters are sampled from the distributions uniform or improper_uniform which both have constant density w.r.t. the Lebesgue measure on their domain. These constants introduce a constant ratio (∝) between the two measures. In addition, since parameters cannot appear on the left-hand side of an assignment we can simplify the return statement. The detailed proof is given in [START_REF] Baudart | Compiling Stan to Generative Probabilistic Languages and Extension to Deep Probabilistic Programming[END_REF].

Score and target. We now show that the value of the Stan target variable corresponds to the score computed by the generated code. Lemma 3.2. For all Stan statements stmt compiled with a continuation k, if γ (target) = 0, and stmt

γ = γ ′ , {[C k (stmt)]} γ = λU . exp(γ ′ (target)) × {[k]} γ ′ [target←0] (U )
The proof is done by induction on the structure of stmt and the finite number of loops iterations (Assumption 1). The hypothesis γ (target) = 0 simplifies the induction by avoiding to keep an accumulator of the value of target. Resetting the value of target in the environment γ ′ [target ← 0] for the evaluation of the continuation k is thus necessary for the inductive step. The proof is given in [START_REF] Baudart | Compiling Stan to Generative Probabilistic Languages and Extension to Deep Probabilistic Programming[END_REF].

Correctness. We now have all the elements to prove that the comprehensive compilation is correct. That is, generated code yields the same un-normalized measure up to a constant factor that will be normalized away by the inference. Theorem 3.3. For all Stan programs p, the semantics of the source and compiled programs are equal up to a constant:

{[p]} D ∝ {[C(p)]} D Proof.
The proof is a direct consequence of Lemmas 3.1 and 3.2 and the definition of the two semantics. With stmt = model(p) and P = params(p):

{[C(p)]} D ∝ λU . ∫ U {[C return(()) (stmt)]} D[P←θ ] ({()}) dθ = λU . ∫ U exp( stmt D[P←θ ] (target)) × {[return(())]}({()}) dθ = λU . ∫ U exp( stmt D[P←θ ] (target)) dθ = {[p]} D □ 4 Implementation
We implemented two new backends for the Stan compiler targeting Pyro [START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF] and NumPyro [START_REF] Phan | Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro[END_REF]. NumPyro is a variant of Pyro built on top of JAX [START_REF] Bradbury | JAX: composable transformations of Python+NumPy programs[END_REF], a Python library that provides efficient automatic differentiation, vectorization, and just-intime compilation on CPU, GPU, and TPU.

For both backends, we have implemented three compilation schemes: generative (Section 2.1), comprehensive (Section 2.3), and mixed.

Mixed compilation. The mixed compilation scheme is an optimization of the comprehensive translation where proper priors are used whenever possible. The mixed compilation can thus generate code that is similar to the generative translation whenever possible.

The mixed translation can be decomposed into three steps: first, compile the program with the comprehensive scheme; second, using the commutativity theorem of [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF], reschedule sample(uniform) statements as late as possible and reschedule observe(D,x) statements as early as possible; and third, merge consecutive sample and observe statements using the following property:

let x = sample(uniform) in let () = observe(D,x) in e ≡ let x = sample(D) in e
In Stan, distribution are automatically truncated based on the parameter support as in the following model. The merge between the sample and observe statements is thus only correct if the two distributions have the same support. We extended the signature of Stan distributions to include the definition domain which can be used to check if the merge is possible.

Finally, the mixed compilation generates correct code even if the ~statements do not respect the dependency order such as in the following example. y ~normal(x, 1); x ~normal(0, 1); ... The mixed compilation reschedules the statements to generate the following code which does not break any environment update: let x = sample(normal(0, 1)) in let y = sample(normal(x, 1)) in ... Architecture. We implemented the compiler as a fork of the Stanc3 compiler 1 , thus guarantying compatibility with the official Stan syntax and existing static analyses. The Stanc3 compiler is composed of multiple intermediate languages. We decided to implement the new backends for the first internal language which is the closest to the Stan source. The implementation is thus closer to the formalization, making it easier to keep track of the correspondence between 1 https://github.com/stan-dev/stanc3 the Stan source and the generated Pyro and NumPyro code. In particular, in Pyro and NumPyro all sampling sites (corresponding to sample and observe) are associated to a unique name which can be used for diagnostics and results analyses. We use Stan variable names with an optional postfix to preserve uniqueness when necessary. For example, in loops, the postfix tracks the current iteration.

Compiling to Pyro. The compiler addresses common challenges like name handling. Pyro and Stan naming conventions are different (e.g., lambda is a common parameter name in Stan) and Pyro has a shared namespace, while Stan distinguishes variables and functions. The compiler carefully avoids conflicts by renaming. Moreover, Stan supports function overloading. The compiler uses static type information to disambiguate function calls by renaming. Finally, there are semantics differences like one-based vs. zero-based arrays.

Stan has a large standard library that also has to be ported to Pyro. Our implementation currently supports a substantial portion of, but not the entire, standard library. Even though Pyro is built on top of Python and thus benefits from a large set of packages, it is not straightforward to implement all Stan functions. The Python counterpart sometimes also has typing or semantic differences. For example, the Stan Bernoulli distribution returns an integer and the Pyro one a float. The differences are handled either in the library or in the compiler. The categorical distribution, which is defined on [1, N ] in Stan and on [0, N -1] in Pyro, illustrates both aspects. The translation of categories is done in the library for a call to the log probability mass function:

def categorical_lpmf(y, theta):

return Categorical(theta).log_prob(y -1)

and the compiler translates (y ~categorical(theta)) into observe(categorical(theta), y -1)

Pyro does not require type declarations, but preserving the shape information of the indexes structures (arrays, vectors, row vectors, matrices) is important. Parameter shapes are passed as arguments to the priors' initialization.

Finally, compared to GProb, Pyro is a Python library. The compiler can thus use Python imperative features and sideeffects for in-place mutation of arrays. However, inference cannot run on models with arbitrary side-effects. For instance, we need to introduce explicit copies when array cells are updated inside loops.

Compiling to NumPyro. The NumPyro backend shares the Pyro backend's challenges and has additional constraints coming from JAX. Dynamic features like dynamic array slices are not supported and will fail during inference. Control structures (conditional and loops) are library functions where the body must be passed in as a pure function. Our compiler accomplishes this by lambda-lifting the bodies of the control structures. This is similar to the compilation of Stan loops to GProb described in Section 3.3 where the updated variables are given explicitly. Returning to the coin example (Figure 1), the compiled NumPyro code using the mixed compilation scheme is thus

def model(N, x): z = sample(beta(1, 1)) def fori__2(i, acc): observe(bernoulli(z), x[i -1]) _ = fori_loop(1, N + 1, fori__2, None)
NumPyro loops are a recent feature which can have a noticeable performance impact. If the body of a loop does not contain probabilistic constructs, we thus generate a JAX loop instead of a NumPyro one.

Extending Stan: Explicit Variational Guides and Neural Networks

Probabilistic languages like Pyro offer new features to program and reason about complex models. This section shows that our compilation scheme can be used to lift these benefits for Stan users. Building on Pyro, we propose DeepStan, a conservative extension of Stan with: (1) variational inference with high-level but explicit guides, and (2) a clean interface to neural networks written in PyTorch.

Explicit Variational Guides

Variational Inference (VI) tries to find the member q θ * (z) of a family Q = q θ (z) θ ∈Θ of simpler distributions that is the closest to the true posterior p(z | x) [START_REF] Blei | Variational Inference: A Review for Statisticians[END_REF]. Members of the family Q are characterized by the values of the variational parameters θ . The fitness of a candidate is measured using the Kullback-Leibler (KL) divergence from the true posterior, which VI aims to minimize:

q θ * (z) = argmin θ ∈Θ KL q θ (z) ∥ p(z | x) .
Pyro natively supports variational inference and lets users define the family Q (the variational guide) alongside the model. To support this for Stan users, we extend Stan with two new optional blocks: guide parameters and guide. The guide block defines a distribution parameterized by the guide parameters. Variational inference optimizes the values of these parameters to approximate the true posterior.

DeepStan inherits restrictions for the definition of the guide from Pyro: the guide must be defined on the same parameter space as the model, i.e., it must sample all the parameters of the model; and the guide should also describe a distribution from which we can directly generate valid samples without running the inference first, which prevents the use of non-generative features and updates of target. random variables for use in the model, guide parameters are learnable coefficients that will be optimized during inference. These restrictions still allow sophisticated guides. The following section presents a guide defined by a neural network.

Adding Neural Networks

One important advantage of Pyro is its tight integration with PyTorch, enabling the authoring of deep probabilistic models: probabilistic models involving neural networks. It is impractical to define neural networks directly in Stan. To support deep probabilistic models, we extend Stan with an optional networks block to import neural network definitions.

Neural networks can be used to capture intricate dynamics between random variables. An example is the Variational Auto-Encoder (VAE) illustrated in Figure 8. A VAE learns a vector-space representation z for each observed data point x (e.g., the pixels of an image) [START_REF] Diederik | Auto-Encoding Variational Bayes[END_REF][START_REF] Jimenez Rezende | Stochastic Backpropagation and Approximate Inference in Deep Generative Models[END_REF]. Each data point x depends on the latent representation z in a complex non-linear way, via a deep neural network: the decoder. The leftmost part of Figure 8 shows the corresponding graphical model. The output of the decoder is a vector µ that parameterizes a Bernoulli distribution over each dimension of x (e.g., each pixel associated to its probability of being in the image).

The key idea of the VAE is to use variational inference to learn the latent representation. The guide maps each x to a latent variable z via another neural network: the encoder. The middle part of Figure 8 shows the graphical model of the guide. The encoder returns, for each input x, the parameters µ z and σ z of a Gaussian distribution in the latent space. Inference tries to learn good values for the parameters θ and ϕ, simultaneously training the decoder and the encoder.

The right part of Figure 8 shows the corresponding code in DeepStan. A network is introduced similarly to an external function with its signature and must be implemented in PyTorch. The network can be used in subsequent blocks, in particular the model block and the guide block.

Bayesian Networks

Neural networks can also be treated as probabilistic models.

A Bayesian neural network is a neural network whose learnable parameters (weights and biases) are random variables instead of concrete values [START_REF] Radford | Bayesian Learning for Neural Networks[END_REF]. Building on Pyro features, we make it easy for users to lift neural networks, i.e., replace concrete neural network parameters by random variables. The left side of Figure 9 shows a simple classifier for handwritten digits based on a multi-layer perceptron (MLP) where all the parameters are lifted to random variables. Unlike the networks used in the VAE, the parameters (regrouped under the variable θ ) are represented using a circle to indicate random variables. The inference starts from prior beliefs about the parameters and learns distributions that fit observed data. We then generate samples of concrete weights and biases to obtain an ensemble of as many MLPs as desired. The ensemble can vote for predictions and can quantify agreement.

The right of Figure 9 shows the corresponding code in DeepStan. We let users declare lifted neural network parameters in Stan's parameters block just like any other random variables. Network parameters are identified by the name of the network and a path, e.g., mlp.l1.weight, following PyTorch naming conventions. The model block defines normal(0,1) priors for the weights and biases of the two linear layers of the MLP. Then, for each image, the computed label follows a categorical distribution parameterized by the output of the network, which associates a probability to each of the ten possible values of the discrete random variable label. The guide parameters define µ and σ , and the guide block uses those parameters to propose normal distributions for the model parameters.

Compiling Bayesian neural networks. To lift neural networks, we use Pyro random_module, a primitive that takes a PyTorch network and a dictionary of distributions and turns the network into a distribution of networks where each parameter is sampled from the corresponding distribution. We treat network parameters as any other random variables and apply the comprehensive translation from Section 2.3. This translation initializes parameters with a uniform prior. It is also possible to mix probabilistic parameters and nonprobabilistic parameters. Our translation only lifts the parameters that are declared in the parameters block by only adding those to the priors dictionary.

Evaluation

We presented our compilation scheme and described the implementation of two new backends for the Stanc3 compiler, targeting Pyro and NumPyro. This section evaluates our compilation scheme and the proposed extensions.

Compiling Stan to Pyro

First we focus on compiling classic Stan models to Pyro and NumPyro. We consider three questions:

RQ1: Can we compile and run all Stan models? RQ2: What is the impact of the compilation on accuracy? RQ3: What is the impact of the compilation on speed? To answer these, we used two publicly available benchmark suites: the example-models2 repository and Posteri-orDB [START_REF] Vehtari | PosteriorDB: a database with data, models and posteriors[END_REF], a database of Stan models with corresponding data, reference posterior samples, and the configuration used to obtain these samples with Stan. The experiments were run on a Linux server with 64 cores (2.10GHz, 40GB RAM) with the latest version of Pyro (1.5.0), NumPyro (0.4.1), and cmdstanpy (0.9.67), without GPUs. The code of the experiments is available at https://github.com/deepppl/evaluation. RQ1: Generality of the compilation. We run our compiler on the 541 models of the example-models repository. Stanc3 semantics checks reject 10 models. Out of the remaining 531, we were able to compile 522 models with the comprehensive and mixed compilation schemes for both the Pyro and NumPyro backends, but only 166 with the generative scheme. This further validates the need of our comprehensive translation. The 9 failures all involve truncations, a feature that is not natively supported in Pyro.

To test the inference, we run 1 iteration on the 98 pairs (models, data) of PosteriorDB that can be compiled with Stanc3. Table 2 presents results for the three compilation schemes: comprehensive, mixed, and generative.

The mixed optimization has no influence on the results. Failures with the Pyro backend are caused by missing standard library functions that are complicated to port to Pyro. As discussed in Section 4, the NumPyro backend relies on JAX, which limits what can be expressed in the model. The additional errors all involve dynamic features that are not supported in JAX.

As a baseline, we run the same experiment with the generative translation. As expected, compilation fails on 60 models. The additional runtime errors are the same as for the comprehensive and mixed translations. RQ2: Accuracy. To evaluate inference accuracy we compare posterior distributions with the criteria used by regression tests for Stan: 3 For each parameter, we check if the error between the means is less than 30% of the standard deviation of the reference. For multidimensional parameters we check the same property for every component:

|mean(θ ref ) -mean(θ )| < 0.3 stddev(θ ref ).
PosteriorDB provides reference samples for 49 pairs (model, dataset). Using Stan with the same configuration (iterations, warmups, chains, thinning, seed), only 31 pairs pass the accuracy test and are thus valid baselines for our evaluation. We run using the Pyro and NumPyro implementation of NUTS with the same configuration 4 , and compare the results with the reference posteriors (NUTS, the No U-Turn Sampler [START_REF] Hoffman | The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo[END_REF], is an optimized HMC and Stan's preferred inference method). Table 3 summarizes the results (additional results on the 49 examples are given in [START_REF] Baudart | Compiling Stan to Generative Probabilistic Languages and Extension to Deep Probabilistic Programming[END_REF]).

Most of the models yield posterior distributions that match the reference. The four remaining errors are due to the following missing functions in our implementation of the standard library: cov_exp_quad (accel_gp and gp_regr) and ODE solvers (lotka_volterra and one_comp_mm_elim_abs). The mismatch (garch11) is due to a constraint that we do not know how to compile in Pyro/NumPyro (the domain of a parameter is constrained by the value of another one).

Table 3 shows that NUTS in Pyro is much slower than its NumPyro counterpart on our examples. Due to the high computational cost of running inference in Pyro (more than 100h for hmm_drive_0), we focus on the NumPyro backend to compare the different compilation schemes. Results show no difference between the comprehensive compilation scheme and the mixed version. When the generative translation is possible, the results also match except for one example (eight_schools_noncentered) due to a parameter constraint that is not propagated in the model (see Section 4).

As discussed in Section 4, the mixed compilation scheme can recover the code produced by the generative translation when possible. Table 3 shows that the extra priors introduced by the translation have no impact on the accuracy of the inference when using NUTS. However, these priors could play a critical role for other inference schemes, e.g., the importance sampling algorithm.

RQ3: Speed. To compare inference speed, Table 3 reports average runtime for Stan and NumPyro over five runs with varying seed values. For obvious computational cost reasons, we only report the duration of one Pyro run. As in RQ2, iterations, warmups, chains, and thinning configurations are given in PosteriorDB.

Table 3 shows that the runtime of DeepStan with the NumPyro backend is competitive with Stan under all three compilation schemes. In addition, runtime durations for the models compiled with the mixed, comprehensive, and generative scheme are almost identical when inference succeeds. These results indicate that the chosen compilation scheme has negligible influence on inference speed. Moreover, as shown in the last column, the NumPyro backend speeds up most benchmarks compared to the highly optimized Stan inference engine (geometric mean: 2.3x on 26 benchmarks).

Two examples (eight_schools-eight_schools_centered and arma11) are very sensitive to random seed variations in Stan and NumPyro (relative standard deviation std/mean > 1). For all other examples std/mean ≤ 0.1. Some of the benchmarks involve nested loops that are still experimental in NumPyro (arK, dogs, dogs_log, hmm_drive_0, hmm_example). For these example NumPyro is typically slower that Stan. If we exclude them we get an overall speedup of 3.8x on 21 benchmarks.

For Table 3 we pre-compiled the models to focus on inference time. The average compilation time for our compiler with both backends was 0.3s (std: 0.02) compared to 10.5s (std: 5.1) for Stan. The NumPyro backend thus outperforms Stan in both compilation and inference speed for these examples.

Stan Extensions

This section evaluates DeepStan, our extension with explicit variational guides and support for deep probabilistic programming with neural networks. We consider two questions: RQ4: Are explicit variational guides useful? RQ5: For deep probabilistic models, how does DeepStan compare to hand-written Pyro code?

RQ4: Explicit guides. The multimodal example shown in Figure 10 is a mixture of two Gaussian distributions with different means but identical variances. The first two histograms of Figure 10 show that in both Stan and DeepStan, this example is particularly challenging for NUTS. Using multiple chains, NUTS finds the two modes, but the chains do not mix and the relative densities are incorrect. This is a known limitation of HMC. 5 Stan also offers ADVI [START_REF] Kucukelbir | Automatic Differentiation Variational Inference[END_REF], an implementation of blackbox VI where guides are automatically synthesized from the model using a full-rank (or mean-field) approximation. 5 https://mc-stan.org/users/documentation/case-studies/identifying_ mixture_models.html These choices imply that ADVI cannot approximate multimodal distribution as illustrated in the last histogram of Figure 10. On the other hand, using the custom variational guide presented in Figure 10, DeepStan with VI is able to find the two modes.

RQ5: Deep probabilistic models. As Stan lacks support for deep probabilistic models, it cannot be used as a baseline. Instead, we compare the performance of the compiled code with hand-written Pyro code on the VAE described in Section 5.2 and a simple Bayesian neural network.

Variational autoencoders were not designed as a predictive model but as a generative model to reconstruct images. Evaluating the performance of a VAE is thus non-obvious. We trained two VAEs on the MNIST dataset using VI: one hand-written in Pyro, the other written in DeepStan. For each image in the test set, the trained VAEs compute a latent representation of dimension 5. We cluster these representations using KMeans with 10 clusters. We measure VAE performance with the pairwise F1 metric: true positives are the number of images of the same digit that appear in the same cluster. For Pyro F1=0.41 (precision=0.43, recall=0.40), and for DeepStan F1=0.43 (precision=0.44, recall=0.42). These numbers shows that compiling DeepStan to Pyro does not impact the performance of such deep probabilistic models. We trained two implementations of a 2-level Bayesian multi-layer perceptron (MLP) with the parameters all lifted to random variables (see section 5.3): one hand-written in Pyro, the other written in DeepStan. We trained both models for 20 epochs on the training set. For each model we generated 100 samples of concrete weights and biases to obtain an ensemble MLP that can be used to compute a distribution of predicted labels. The accuracy for both models is 92% on the test set and the agreement between the two models is above 95%. The execution time is comparable. These experiments show that compiling DeepStan models to Pyro has little impact on the model. Changing the priors on the network parameters from normal(0,1) to normal(0,10) (see Section 5.3) increases accuracy from 0.92 to 0.96. This further validates our compilation, compiling parameter priors to observe statements on deep probabilistic models.

Related Work

To the best of our knowledge, we propose the first comprehensive translation of Stan to a generative PPL. The closest related work was developed by the Pyro team [START_REF] Chen | Transpiling Stan models to Pyro[END_REF]. Their work focuses on performance and ours on completeness. Their proposed compilation technique corresponds to the generative translation presented in Section 2.1 and thus only handles a subset of Stan. The code is not open-source, and we rely on our own implementation of the generative translation in Section 6. Compared to our approach, they are also looking into independence assumptions between loop iterations to generate parallel code. Combining these ideas with our approach is a promising future direction. They do not extend Stan with either VI or neural networks. Similarly, in Appendix B.2 of [START_REF] Gorinova | Probabilistic programming with densities in SlicStan: efficient, flexible, and deterministic[END_REF], Gorinova et al. outline the generative translation of Section 2.1, and also mention the issue with multiple updates but do not provide a solution. Lee et al. [START_REF] Lee | Towards verified stochastic variational inference for probabilistic programs[END_REF] introduce a density-based semantics for Pyro, but this semantics does not handle Stan's non-generative features.

The goal of compiling Stan to Pyro is to create a platform for experimenting with new ideas. For example, Section 5.1 extends Stan with explicit variational guides. Similarly, Pyro now offers inference on discrete parameters 6 that we could port to Stan using our backends.

In recent years, taking advantage of the maturity of DL frameworks, multiple deep probabilistic programming languages have been proposed: PyMC3 [START_REF] Salvatier | Probabilistic programming in Python using PyMC3[END_REF] built on top of Theano, Edward [START_REF] Tran | Deep Probabilistic Programming[END_REF] and ZhuSuan [START_REF] Shi | ZhuSuan: A Library for Bayesian Deep Learning[END_REF] built on top of Ten-sorFlow, and Pyro [START_REF] Bingham | Pyro: Deep Universal Probabilistic Programming[END_REF] and ProbTorch [START_REF] Narayanaswamy | Learning Disentangled Representations with Semi-Supervised Deep Generative Models[END_REF] built on top of PyTorch. All these languages are implemented as libraries. The users thus need to master the entire technology stack of the library, the underlying DL framework, and the host language. In comparison, DeepStan is a self-contained language and the compiler helps the programmer via dedicated static analyses.

Conclusion

This paper introduces a comprehensive compilation scheme from Stan to generative probabilistic programming languages. This shows that Stan is at most as expressive as this family of languages. We implemented a compiler from Stan to Pyro. Additionally, we designed and implemented extensions for Stan with explicit variational guides and neural networks. where for all 0 ≤ 𝑖 < 𝑛, 𝛾 𝑖 = 𝑠 𝛾 𝑖+1 , and 𝛾 𝑛 = 𝛾.

Loops. By

If 𝑛 = 0, 𝑒 𝛾 0 [target←𝑡 ] = 𝑒 𝛾 0 [target←0] = 0. So by definition of the semantics we have:

𝛾 0 [target ← 𝑡] (target) = 𝑡 + 𝛾 0 [target ← 0] (target)
For the inductive case, if 𝑛 = 𝑖 +1, there are some iterations left so by definition of the semantics we have:

{[C(𝑝)]} 𝛾 ∝ 𝜆𝑈 . ∫ 𝑈 {[C return(()) (stmt)]} 𝛾 [ P←𝜃 ] ({()})𝑑𝜃
Proof. Let P = 𝑥 1 , ..., 𝑥 𝑛 . By definition of the compilation function, C(𝑝) has the shape:

let 𝑥 1 = 𝐷 1 in . . . let 𝑥 𝑛 = 𝐷 𝑛 in C return(P) (stmt)
where for each parameter 𝑥 𝑖 , the distribution 𝐷 𝑖 is either uniform or improper_uniform. In both cases the corresponding density is constant w.r.t. the Lebesgue measure on its domain.

Since the kernels defined by GProb expressions are always s-finite [START_REF] Staton | Commutative Semantics for Probabilistic Programming[END_REF], from the semantics of GProb (Section 3.2) and the Fubini-Tonelli theorem, we have: Sequence. If 𝛾 1 = stmt 1 𝛾 and 𝛾 2 = stmt 2 𝛾 1 [target←0] , the induction hypothesis and the semantics of GProb yield:

{[C(𝑝)]} 𝛾 = {[let 𝑥 1 = 𝐷 1 in . . . let 𝑥 𝑛 = 𝐷 𝑛 in C return(P) (stmt)]} 𝛾 = 𝜆𝑈 . ∫ 𝑋 1 𝐷 1 (
{[C 𝑘 (stmt 1 ; stmt 2 )]} 𝛾 = {[C C 𝑘 (stmt 2 ) (stmt 1 )]} 𝛾 = { by induction on stmt 1 and definition of 𝛾 1 } 𝜆𝑈 .exp(𝛾 1 (target)) × {[C 𝑘 (stmt 2 )]} 𝛾 1 [target←0] (𝑈 ) = { by induction on stmt 2 and definition of 𝛾 2 } 𝜆𝑈 .exp(𝛾 1 (target)) × exp(𝛾 2 (target)) × {[𝑘]} 𝛾 2 [target←0] (𝑈 ) = 𝜆𝑈 .exp(𝛾 1 (target) + 𝛾 2 (target)) × {[𝑘]} 𝛾 2 [target←0] (𝑈 )
On the other hand, from Lemma A.2 we have: stmt 1 ; stmt 2 𝛾 (target) = 𝛾 1 (target) +𝛾 2 (target) which conclude the proof of this case.

Loop. For the case while (𝑒) {𝑠}, we first look at the semantics of the loop on the Stan side. We can prove that

while (𝑒) {𝑠} 𝛾 (target) = 𝑛-1 𝑖=0 𝛾 𝑖 (target) (1)
where 𝑛 is the number of loop iterations, for all 0 ≤ 𝑖 < 𝑛, 𝛾 𝑖 = 𝑠 𝛾 𝑖+1 [target←0] , and 𝛾 𝑛 = 𝛾. By Assumption 1 the programs terminate, so 𝑛 is finite. This is a direct consequence of the following property since 𝛾 (target) = 0:

while (𝑒) {𝑠} 𝛾 𝑛 [target←0] (target) = 𝑛-1 𝑖=0 𝛾 𝑖 (target) (2)
The proof of eq. ( 2) is done by induction on the number of remaining iterations. If 𝑛 = 0, the condition of the loop is false. Since the condition of the loop cannot depend on the value of target, it means that 𝑒 𝛾 0 = 𝑒 𝛾 0 [target←0] = 0, so by application of the semantics we have directly:

while (𝑒) {𝑠} 𝛾 0 [target←0] (target) = 0
For the inductive case (𝑛 = 𝑖 + 1), there is some loop iterations remaining and thus 𝑒 𝛾 𝑖+1 = 𝑒 𝛾 𝑖+1 [target←0] ≠ 0. By application on the semantics we have: Let's look at the loop while X (𝑒) C return(X) (𝑠) where X is the variables updated by 𝑠 (lhs(𝑠) = X). We will use the two following properties. For all 𝛾, if 𝑒 𝛾 = 0

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 = 𝛿 𝛾 ( X) (3)
For all 𝛾, if 𝑒 𝛾 ≠ 0 and 𝑠 𝛾 = 𝛾 ′ :

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 = 𝜆𝑈 .exp(𝛾 ′ (target)) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 ′ [target←0] (𝑈 ) (4)
The proofs of both of these properties start by unfolding the definition of the semantics:

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 = 𝜆𝑈 .if 𝑒 𝛾 = 0 then 𝛿 𝛾 ( X) (𝑈 ) else ∫ 𝑋 {[C return(X) (𝑠)]} 𝛾 (𝑑𝑣) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 [X←𝑣 ] (𝑈 )
From there, the proof of eq. ( 3) where 𝑒 𝛾 = 0 is trivial. For eq. ( 4) where 𝑒 𝛾 ≠ 0 and 𝑠 𝛾 = 𝛾 ′ , we apply the induction

hypothesis on {[C return(X) (𝑠)]} 𝛾 : {[while X (𝑒) C return(X) (𝑠)]} 𝛾 = { by simplification of 𝑒 𝛾 ≠ 0} 𝜆𝑈 . ∫ 𝑋 {[C return(X) (𝑠)]} 𝛾 (𝑑𝑣) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 [X←𝑣 ] (𝑈 )
= { by induction on the structure of the statement} 𝜆𝑈 .

∫ 𝑋 (exp(𝛾 ′ (target)) × {[return(X)]} 𝛾 ′ [target←0] )(𝑑𝑣) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 [X←𝑣 ] (𝑈 ) = { by definition of the semantics of return} 𝜆𝑈 . ∫ 𝑋 (exp(𝛾 ′ (target)) × 𝛿 𝛾 ′ [target←0] ( X) )(𝑑𝑣) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 [X←𝑣 ] (𝑈 ) = { by the integration of the 𝛿 distribution } 𝜆𝑈 .exp(𝛾 ′ (target)) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 [X←𝛾 ′ [target←0] ( X) ] (𝑈 )
= { the only difference between 𝛾 and 𝛾 ′ are in X by def of lhs }

𝜆𝑈 .exp(𝛾 ′ (target)) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 ′ [target←0] (𝑈 )
We can now prove by induction on the number of loop iterations 𝑛 (that we assume finite) that

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 𝑛 [target←0] = 𝜆𝑈 . 𝑛-1 𝑖=0 exp(𝛾 𝑖 (target)) × 𝛿 𝛾 0 ( X) (𝑈 ) (5)
where for all 0 ≤ 𝑖 < 𝑛, 𝛾 𝑖 = 𝑠 𝛾 𝑖+1 [target←0] .

In the base case of the induction (𝑛 = 0), there is no more iteration ( 𝑒 𝛾 0 = 𝑒 𝛾 0 [target←0] = 0). By application of eq. ( 3), we have:

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 0 [target←0] = 𝛿 𝛾 ( X)
For the inductive case (𝑛 = 𝑖 + 1), we know that 𝑒 𝛾 𝑖+1 = 𝑒 𝛾 𝑖+1 [target←0] ≠ 0 since there are some more loop iterations to execute. Therefore, we can apply eq. ( 4):

{[while X (𝑒) C return(X) (𝑠)]} 𝛾 𝑖+1 [target←0]
= { by eq. ( 4) and definition of 𝛾 𝑖 }

𝜆𝑈 .exp(𝛾

𝑖 (target)) × {[while X (𝑒) C return(X) (𝑠)]} 𝛾 𝑖 [target←0] (𝑈 ) = { by induction on the number of iterations } 𝜆𝑈 .exp(𝛾 𝑖 (target)) × 𝑖-1 𝑗=0 exp(𝛾 𝑗 (target)) × 𝛿 𝛾 0 ( X) (𝑈 ) = 𝜆𝑈 . 𝑖 𝑗=0 exp(𝛾 𝑗 (target)) × 𝛿 𝛾 0 ( X) (𝑈 )
We can finally put all the pieces together to prove the case of while (𝑒) {𝑠}. We still assume that the loop has 𝑛 iterations. We first unroll the definition of compilation function and the semantics.

{[C 𝑘 (while (𝑒) {𝑠})]} 𝛾 = {[let X = while X (𝑒) C return(X) (𝑠) in 𝑘]} 𝛾 = 𝜆𝑈 . ∫ 𝑋 {[while X (𝑒) C return(X) (𝑠)]} 𝛾 (𝑑𝑣) × {[𝑘]} 𝛾 [X←𝑣 ] (𝑈 )
= { by eq. ( 5) with 𝛾 𝑛 = 𝛾 and because 𝛾 (target) = 0 } 𝜆𝑈 .

∫ 𝑋 𝑛-1 𝑖=0 exp(𝛾 𝑖 (target)) × 𝛿 𝛾 0 ( X) (𝑑𝑣) × {[𝑘]} 𝛾 [X←𝑣 ] (𝑈 ) = { by rewriting the product of exponential } 𝜆𝑈 . ∫ 𝑋 exp( 𝑛-1 𝑖=0 𝛾 𝑖 (target)) × 𝛿 𝛾 0 ( X) (𝑑𝑣) × {[𝑘]} 𝛾 [X←𝑣 ] (𝑈 ) = { by eq. (1) with 𝛾 ′ = while (𝑒) {𝑠} 𝛾 } 𝜆𝑈 . ∫ 𝑋 exp(𝛾 ′ (target)) × 𝛿 𝛾 0 ( X) (𝑑𝑣) × {[𝑘]} 𝛾 [X←𝑣 ] (𝑈 ) = { by the integration of the 𝛿 distribution } 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 [X←𝛾 0 ( X) ] (𝑈 ) = { by definition of lhs and since target ∉ X } 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 0 [target←0] (𝑈 ) = 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 ′ [target←0] (𝑈 )

□

C Evaluation

We present here the results of the experiments for RQ2 and RQ3 described in Section 6 on the entire set of PosteriorDB models with reference, that is, without filtering out the models for which Stan does not pass the accuracy test. RQ2: Accuracy. Table 4 presents the accuracy results. For each run we report the mean relative error in green if the model pass the accuracy test, orange otherwise (the accuracy test fails if the relative error for one of the parameter is above 0.3).

Compared to the results presented in Section 6 we have two additional errors for the NumPyro backend: gp_pois_regr (missing cov_exp_quad) and diamonds_diamonds (missing student_t_lccdf). We also have two additional mismatches low_dim_gauss_mix and hmm_drive_1 both due to an ordered constraint that is not supported in the versions of Pyro and NumPyro that we used.

For the Pyro backend two examples for which Stan returns a mismatch timed-out after 10h (indicated with a "-" in Table 4): hmm_drive_1 and low_dim_gauss_mix.

RQ3: Speed. Table 5 presents the speed results. For each run we report the average duration in second and the standard deviation.

Figure 1 .

 1 Figure 1. Biased coin model in Stan.

Figure 2 .

 2 Figure 2. Compiled coin model of Figure 1.

  decl ::= base_type constraint x ; | base_type constraint x [shape] ; base_type ::= real | int | vector[size] | matrix[size,size] constraint ::= ε | < lower = e , upper = e > | < lower = e > | < upper = e > shape ::= size | shape , size size ::= e Inside a block, Stan is similar to a classic imperative language, with two extra, specialized statements: target += e directly updates the log-density of the model (stored in the reserved variable target), and x ~D indicates that a variable x follows a distribution D.

  For a Stan program p = data {decls d } parameters {decls p } model {decls m stmt}: data(p) = decls d params(p) = decls p model(p) = stmt

Figure 4 .

 4 Figure 4. Semantics of expressions

e

  ::= c | x | {e 1 ,...,e n } | [e 1 ,...,e n ] | e 1 [e 2 ] | f (e 1 ,...,e n ) | let x = e 1 in e 2 | let x[e 1 ,...,e n ] = e in e ′ | if (e) e 1 else e 2 | for X (x in e 1 :e 2 ) e 3 | while X (e 1 ) e 2 | factor(e) | sample(e) | return(e)

Figure 7

 7 defines compilation for statements, C k (stmt), parameterized by a continuation k. Stan imperative assignments become functional updates using local bindings. Compiling the sequences chains the continuations. Updates to the target are compiled into factor expressions and all ~statements are compiled into observations.

  parameters { real<lower=0> sigma; } model { sigma ~normal(0, 1); }

Figure 8 .

 8 Figure 8. Graphical models and DeepStan code of the Variational Auto-Encoder model and guide.

Figure 9 .

 9 Figure 9. Graphical models and DeepStan code of the Bayesian MLP.

  priors = {} priors['l1.weight'] = improper_uniform(shape=[nh, nx]) ... # priors of the other parameters lifted_mlp = pyro.random_module('mlp', mlp, priors)() Then, the Stan ~statements in the model block are compiled into Pyro observe statements. mlp_params = dict(lifted_mlp.named_parameters()) observe(normal(0, 1), mlp_params['l1.weight'])

Figure 10 .

 10 Figure 10. DeepStan code and histograms of the multimodal example using Stan, DeepStan with NUTS, DeepStan with VI, and Stan with ADVI.

  while (𝑒) {𝑠} 𝛾 𝑖+1 [target←𝑡 ] (target) = while (𝑒) {𝑠} 𝑠 𝛾 𝑖+1 [target←𝑡 ] (target) { by definition } = { by definition of 𝛾 𝑖 and Lemma A.1 } while (𝑒) {𝑠} 𝛾 𝑖 [target← 𝑠 𝛾 𝑖+1 [target←t] (target) ] (target) = { by induction on the number of reductions } 𝑠 𝛾 𝑖+1 [target←t] (target) + while (𝑒) {𝑠} 𝛾 𝑖 [target←0] (target) = { by induction on 𝑠 } 𝑡 + 𝑠 𝛾 𝑖+1 [target←0] (target) + while (𝑒) {𝑠} 𝛾 𝑖 [target←0] (target) = { by definition of 𝛾 𝑖 } 𝑡 + 𝑠 𝛾 𝑖+1 [target←0] (target) + while (𝑒) {𝑠} 𝑠 𝛾 𝑖+1 [target←0] (target) = { by Lemma A.1 } 𝑡 + 𝑠 𝛾 𝑖+1 [target←0] (target) + while (𝑒) {𝑠} 𝑠 𝛾 𝑖+1 [target←0] [target←0] (target) = { by induction } 𝑡 + while (𝑒) {𝑠} 𝛾 𝑖+1 [target←0] (target) □ B Proofs: Correctness of the Compilation Lemma B.1. For all Stan programs 𝑝 with stmt = model(𝑝) and P = params(𝑝), and environments 𝛾:

{

  [C return(P) (stmt)]} 𝛾 [ P←𝜃 ] (𝑈 ) = ∫ 𝑋 {[C return(()) (stmt)]} 𝛾 [ P←𝜃 ] (𝑑𝑥) × {[return(P)]} [ P←𝜃 ] (𝑈 ) = {[C return(()) (stmt)]} 𝛾 [ P←𝜃 ] ({()}) × 𝛿 𝜃 (𝑈 )Going back to the previous equation, we now have{[C(𝑝)]} 𝛾 = 𝜆𝑈 . ∫ 𝑋 {[C return(()) (stmt)]} 𝛾 [ P←𝜃 ] ({()}) × 𝛿 𝜃 (𝑈 ) 𝑑𝜃 = 𝜆𝑈 . ∫ 𝑈 {[C return(()) (stmt)]} 𝛾 [ P←𝜃 ] ({()}) 𝑑𝜃 □ Lemma B.2.For all Stan statements stmt compiled with a continuation 𝑘, if 𝛾 (target) = 0, and stmt𝛾 = 𝛾 ′ , {[C 𝑘 (stmt)]} 𝛾 = 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 ′ [target←0] (𝑈 )Proof. The proof is done by induction on the structure of stmt and number of reductions using the definition of the compilation function (Section 3.3) and the semantics of GProb.Assignment. Evaluating 𝑥 = 𝑒 does not update target and its initial value is 0 by hypothesis. With 𝛾 ′ = 𝑥 = 𝑒 𝛾 we have 𝛾 ′ [target ← 0] = 𝛾 ′ and exp(𝛾 ′ (target)) = 1. Then from GProb's semantics we have:

{

  [C 𝑘 (𝑥 = 𝑒)]} 𝛾 = { by definition of C 𝑘 (.) } {[let 𝑥 = return(𝑒) in 𝑘]} 𝛾 = { by definition of the semantics } 𝜆𝑈 . ∫ 𝑋 {[return(𝑒)]} 𝛾 (𝑑𝑣) × {[𝑘]} 𝛾 [𝑥←𝑣 ] (𝑈 ) = { by definition of the semantics } 𝜆𝑈 .∫ 𝑋 𝛿 𝑒 𝛾 (𝑑𝑣) × {[𝑘]} 𝛾 [𝑥←𝑣 ] (𝑈 ) = { by the integration of the 𝛿 distribution } 𝜆𝑈 .1 × {[𝑘]} 𝛾 [𝑥← 𝑒 𝛾 ] (𝑈 ) = { by the semantics of 𝑥 = 𝑒 𝛾 } 𝜆𝑈 .1 × {[𝑘]} 𝑥 = 𝑒 𝛾 (𝑈 ) = { by definition of 𝛾 ′ } 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 ′ [target←0] (𝑈 )Target update. Since the evaluation of target += 𝑒 only updates the value of target and its initial value is 0, with 𝛾 ′ = target += 𝑒 𝛾 we have 𝛾 = 𝛾 ′ [target ← 0], and 𝛾 ′ (target) = 𝑒 𝛾 . Then from GProb's semantics we have:{[C 𝑘 (target += 𝑒)]} 𝛾 = { by definition of C 𝑘 (.) } {[let () = factor(𝑒) in 𝑘]} 𝛾 = { by definition of the semantics } 𝜆𝑈 . ∫ () exp( 𝑒 𝛾 )𝛿 () (𝑑𝑣) × {[𝑘]} 𝛾 (𝑈 ) = {by the integration of the 𝛿 distribution } 𝜆𝑈 . exp( 𝑒 𝛾 ) × {[𝑘]} 𝛾 (𝑈 ) = { by the semantics of target += 𝑒 𝛾 } 𝜆𝑈 . exp( target += 𝑒 𝛾 (target)) × {[𝑘]} 𝛾 (𝑈 ) = { by definition of 𝛾 ′ } 𝜆𝑈 . exp(𝛾 ′ (target)) × {[𝑘]} 𝛾 ′ [target←0] (𝑈 )

1 𝑗=0𝛾

 1 while (𝑒) {𝑠} 𝛾 𝑖+1 [target←0] (target) = while (𝑒) {𝑠} 𝑠 𝛾 𝑖+1 [target←0] = while (𝑒) {𝑠} 𝛾 𝑖 { by definition of 𝛾 𝑖 } = 𝛾 𝑖 + while (𝑒) {𝑠} 𝛾 𝑖 [target←0] { by Lemma A.2 } = { by induction on the number of iterations } 𝛾 𝑖 + 𝑖-𝑗 (target) = 𝑖 𝑗=0 𝛾 𝑗 (target)

  [e 1 ,...,e n ] = e γ = γ [x ← (x[ e 1 γ ,..., e n γ ] ← e γ )] s 1 ; s 2 γ = s 2 s 1 γ for (x in e 1 :e 2 ) {s} γ = let n 1 = e 1 γ in let n 2 = e 2 γ in if n 1 > n 2 then γ else for (x in n 1 + 1:n 2 ) {s} s γ [x ←n 1 ]

	while (e) {s} γ	= if e γ = 0 then γ else while (e) {s} s γ
	if (e) s 1 else s 2 γ = if e 1 γ 0 then s 1 γ else s 2 γ
	skip γ	= γ
	target += e γ	= γ [target ← γ (target) + e γ ]
	e 1 ~e2 γ	= let D = e 2 γ in target += D lpdf (e 1 ) γ
	Figure 3. Semantics of statements
	Given the input data, the posterior distribution of a Stan pro-
	gram is obtained by normalizing the measure {[p]} D . As the
	integrals are often intractable, the runtime uses approximate
	inference schemes to compute the posterior distribution.

  An if statement executes one of the two branches depending on the value of the condition. A skip leaves the environment unchanged. A statement target += e adds the value of e to target in the environment. Finally, a c γ = c {e 1 ,...,e n } γ = { e 1 γ ,..., e n γ } x γ = γ (x) [e 1 ,...,e n ] γ = [ e 1 γ ,..., e n γ ] e 1 [e 2 ] γ = e 1 γ [ e 2 γ ] f (e) γ = f ( e γ )

  The generative translation from Section 2.1 generates a Python function that can serve as a Pyro guide. The guide parameters block is used to generate Pyro param statements, which introduce learnable parameters. Unlike Stan parameters that define

					networks {
					real[,] decoder(real[] x);
					real[,] encoder(int[,] x); }
		z	z		data { int nz;
					int<lower=0, upper=1> x[28, 28]; }
	θ	decoder	Normal		parameters {
		µ	µz, σz		real z[nz]; }
					model {
		Bernoulli	encoder	ϕ	real mu[28, 28];
					z ~normal(0, 1);
		x	x		mu = decoder(z); x ~bernoulli(mu); }
		N	N		guide {
		model p θ (x | z)	guide q ϕ (z | x)	real encoded[2, nz] = encoder(x);
					real mu_z[nz] = encoded[1];
					real sigma_z[nz] = encoded[2];
					z ~normal(mu_z, sigma_z); }

Table 2 .

 2 Successful inference run for 98 PosteriorDB models.

		Compr. Mixed Gener.
	Pyro	87	87	36
	NumPyro	83	83	35

Table 3 .

 3 Comparing inference results with PosteriorDB references

	Pyro	NumPyro

✓ match, ❍ mismatch, ✗ error. Durations are reported in hh:mm:ss format. Speedup = Stan / NumPyro Compr.

  Assumption 1, there is a finite number 𝑛 of loop iterations. So we can prove by induction on the number of remaining iterations that while (𝑒) {𝑠} 𝛾 𝑛 [target←𝑡 ] (target) = 𝑡 + while (𝑒) {𝑠} 𝛾 𝑛 [target←0] (target)

  𝑑𝑣 1 ) ... ∫ 𝛾 [ (𝑥 1 ,...,𝑥 𝑛 )←(𝑣 1 ,...,𝑣 𝑛 ) ] (𝑈 ) 𝛾 [ (𝑥 1 ,...,𝑥 𝑛 )←(𝑣 1 ,...,𝑣 𝑛 ) ] (𝑈 ) 𝑑𝑣 1 ...𝑑𝑣 𝑛 𝛾 [ (𝑥 1 ,...,𝑥 𝑛 )←𝜃 ] (𝑈 ) 𝑑𝜃 The evaluation of {[C return(P) (stmt)]} 𝛾 [ P←𝜃 ] (𝑈 ) terminates with the value {[return(P)]} 𝛾 ′ [ P←𝜃 ] (𝑈 ) where 𝛾 ′ is the environment obtained after the evaluation of the model statements stmt.However, in Stan, parameters declared in P cannot appear in the left-hand side of an assignment. The evaluation of the model statements stmt cannot update the value of the parameters. We can thus simplify {[return(P)]} 𝛾 ′ [ P←𝜃 ] (𝑈 ) into {[return(P)]} [ P←𝜃 ] (𝑈 ). The evaluation of the compiled code can then be decomposed as follows:

	𝑋 𝑛 𝐷 𝑛 (𝑑𝑣 𝑛 )
	∫ {[C return(P) (stmt)]} ∝ 𝜆𝑈 . ∫ 𝑋 1 ... 𝑋 𝑛 ∫ {[C return(P) (stmt)]} = 𝜆𝑈 . 𝑋 =𝑋 1 ×...×𝑋 𝑛 {[C return(P) (stmt)]}

Table 4 .

 4 Comparing inference results with PosteriorDB references. For each column we report the mean relative error in green if the example pass the accuracy test, orange otherwise.

	Pyro	NumPyro

Table 5 .

 5 Comparing inference speed on PosteriorDB models. For each columns we report the duration in second with the following format mean(std). Results are averaged over 5 runs.

	NumPyro

https://github.com/stan-dev/example-models

https://github.com/stan-dev/performance-tests-cmdstan

The configuration interface for NUTS in Pyro and NumPyro is similar to the CmdStanPy interface.

https://pyro.ai/examples/enumeration.html

Acknowledgements

The authors are greatful to the following people for their helpful feedback and encouragements: E. Bingham, K. Kate, Y. Mroueh, F. Obermeyer, A. Pauthier, and D. Phan.

A Auxiliary Lemmas

The first additional property that we need is that the initial value of target only impact the final value of target but no other variable.

Lemma A.1. For all Stan statement stmt, for all variable 𝑥 different from target, then for all values 𝑡 1 and 𝑡 2 ,

Proof. The proof is done by induction on the structure of the statements and number of reductions. We detail just the most interesting cases.

Assignment. Evaluating 𝑥 = 𝑒 does not update target and according to Assumption 2, the expression 𝑒 does not depends on the value of target.

Target Update. Since the evaluation of target += 𝑒 only updates the value of target, we have by definition

Sequence. For a sequence stmt 1 ; stmt 2 we have:

Loop. For a loop while (𝑒) {𝑠}, by induction on the structure of the statements, for all environment 𝛾, values 𝑡 1 and 𝑡 2 and variable 𝑥 ≠ target, we have 𝑠 𝛾 [target←𝑡 1 ] (𝑥) = 𝑠 𝛾 [target←𝑡 2 ] (𝑥), which implies:

In addition, by Assumption 2 we know that for all environment 𝛾, 𝑒

By Assumption 1, there is a finite number 𝑛 of loop iterations. So we can prove by induction on the number of remaining iterations that for all 𝑡 1 , 𝑡 2 :

where for all 0 ≤ 𝑖 < 𝑛, 𝛾 𝑖 = 𝑠 𝛾 𝑖+1 , and 𝛾 𝑛 = 𝛾.

So by definition of the semantics we have:

For the inductive case, if 𝑛 = 𝑖 + 1, there are some iterations left so 𝑒 𝛾 𝑖+1 [target←𝑡 1 ] = 𝑒 𝛾 𝑖+1 [target←𝑡 2 ] ≠ 0. So by definition of the semantics we have:

□

The next lemma state that every statement can be evaluated in an environment where target is set to zero.

Lemma A.2. For all Stan statement stmt and any real value 𝑡 we have:

Proof. The proof is done by induction on the structure of the statements and number of reductions.