TuringLang · mhauru · Aug 21, 2024 · Oct 31, 2023 · Nov 7, 2023 · Nov 7, 2023
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -4,6 +4,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -14,6 +15,7 @@ DataStructures = "0.18"
 Distributions = "0.25"
 Documenter = "1"
 FillArrays = "0.13, 1"
+ForwardDiff = "0.10"
 LogDensityProblems = "2"
 MCMCChains = "5, 6"
 StableRNGs = "1"
diff --git a/docs/make.jl b/docs/make.jl
@@ -20,8 +20,10 @@ makedocs(;
         "Home" => "index.md",
         "API" => "api.md",
         "Tutorials" => ["tutorials/prob-interface.md"],
+        "Internals" => ["internals/transformations.md"],
     ],
     checkdocs=:exports,
+    doctest=false,
 )
 
 deploydocs(; repo="github.com/TuringLang/DynamicPPL.jl.git", push_preview=true)
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -218,13 +218,62 @@ Please see the documentation of [AbstractPPL.jl](https://github.com/TuringLang/A
 
 ### Data Structures of Variables
 
-DynamicPPL provides different data structures for samples from the model and their log density.
-All of them are subtypes of [`AbstractVarInfo`](@ref).
+DynamicPPL provides different data structures used in for storing samples and accumulation of the log-probabilities, all of which are subtypes of [`AbstractVarInfo`](@ref).
 
 ```@docs
 AbstractVarInfo
 ```
 
+But exactly how a [`AbstractVarInfo`](@ref) stores this information can vary.
+
+#### `VarInfo`
+
+```@docs
+VarInfo
+TypedVarInfo
+```
+
+One main characteristic of [`VarInfo`](@ref) is that samples are stored in a linearized form.
+
+```@docs
+link!
+invlink!
+```
+
+```@docs
+set_flag!
+unset_flag!
+is_flagged
+```
+
+For Gibbs sampling the following functions were added.
+
+```@docs
+setgid!
+updategid!
+```
+
+The following functions were used for sequential Monte Carlo methods.
+
+```@docs
+get_num_produce
+set_num_produce!
+increment_num_produce!
+reset_num_produce!
+setorder!
+set_retained_vns_del_by_spl!
+```
+
+```@docs
+Base.empty!
+```
+
+#### `SimpleVarInfo`
+
+```@docs
+SimpleVarInfo
+```
+
 ### Common API
 
 #### Accumulation of log-probabilities
@@ -241,7 +290,7 @@ resetlogp!!
 ```@docs
 keys
 getindex
-DynamicPPL.getindex_raw
+DynamicPPL.getindex_internal
 push!!
 empty!!
 isempty
@@ -269,8 +318,9 @@ DynamicPPL.invlink
 DynamicPPL.link!!
 DynamicPPL.invlink!!
 DynamicPPL.default_transformation
+DynamicPPL.link_transform
+DynamicPPL.invlink_transform
 DynamicPPL.maybe_invlink_before_eval!!
-DynamicPPL.reconstruct
 ```
 
 #### Utils
@@ -283,56 +333,6 @@ DynamicPPL.varname_leaves
 DynamicPPL.varname_and_value_leaves
 ```
 
-#### `SimpleVarInfo`
-
-```@docs
-SimpleVarInfo
-```
-
-#### `VarInfo`
-
-Another data structure is [`VarInfo`](@ref).
-
-```@docs
-VarInfo
-TypedVarInfo
-```
-
-One main characteristic of [`VarInfo`](@ref) is that samples are stored in a linearized form.
-
-```@docs
-link!
-invlink!
-```
-
-```@docs
-set_flag!
-unset_flag!
-is_flagged
-```
-
-For Gibbs sampling the following functions were added.
-
-```@docs
-setgid!
-updategid!
-```
-
-The following functions were used for sequential Monte Carlo methods.
-
-```@docs
-get_num_produce
-set_num_produce!
-increment_num_produce!
-reset_num_produce!
-setorder!
-set_retained_vns_del_by_spl!
-```
-
-```@docs
-Base.empty!
-```
-
 ### Evaluation Contexts
 
 Internally, both sampling and evaluation of log densities are performed with [`AbstractPPL.evaluate!!`](@ref).

diff --git a/docs/src/assets/images/transformations-assume-without-istrans.dot b/docs/src/assets/images/transformations-assume-without-istrans.dot
@@ -0,0 +1,17 @@
+digraph {
+  # `assume` block
+  subgraph cluster_assume {
+    label = "assume";
+    fontname = "Courier";
+
+    assume [shape=box, label=< assume(varinfo, <FONT COLOR="#3B6EA8">@varname</FONT>(x), Normal())>, fontname="Courier"];
+    without_linking_assume [shape=box, label="f = from_internal_transform(varinfo, varname, dist)", fontname="Courier"];
+    with_logabsdetjac [shape=box, label="x, logjac = with_logabsdet_jacobian(f, assume_internal(varinfo, varname, dist))", fontname="Courier"];
+    return_assume [shape=box, label=< <FONT COLOR="#3B6EA8">return</FONT> x, logpdf(dist, x) - logjac, varinfo >, style=dashed, fontname="Courier"];
+
+    assume -> without_linking_assume;
+    without_linking_assume -> with_logabsdetjac;
+    with_logabsdetjac -> return_assume;
+  }
+}
+
diff --git a/docs/src/assets/images/transformations-assume-without-istrans.dot.png b/docs/src/assets/images/transformations-assume-without-istrans.dot.png
diff --git a/docs/src/assets/images/transformations-assume-without-istrans.dot.svg b/docs/src/assets/images/transformations-assume-without-istrans.dot.svg
diff --git a/docs/src/assets/images/transformations-assume.dot b/docs/src/assets/images/transformations-assume.dot
@@ -0,0 +1,22 @@
+digraph {  
+  # `assume` block
+  subgraph cluster_assume {
+    label = "assume";
+    fontname = "Courier";
+
+    assume [shape=box, label=< assume(varinfo, <FONT COLOR="#3B6EA8">@varname</FONT>(x), Normal())>, fontname="Courier"];
+    iflinked_assume [label=< <FONT COLOR="#3B6EA8">if</FONT> istrans(varinfo, varname) >, fontname="Courier"];
+    without_linking_assume [shape=box, label="f = from_internal_transform(varinfo, varname, dist)", fontname="Courier"];
+    with_linking_assume [shape=box, label="f = from_linked_internal_transform(varinfo, varname, dist)", fontname="Courier"];
+    with_logabsdetjac [shape=box, label="x, logjac = with_logabsdet_jacobian(f, assume_internal(varinfo, varname, dist))", fontname="Courier"];
+    return_assume [shape=box, label=< <FONT COLOR="#3B6EA8">return</FONT> x, logpdf(dist, x) - logjac, varinfo >, style=dashed, fontname="Courier"];
+
+    assume -> iflinked_assume;
+    iflinked_assume -> without_linking_assume [label=<  <FONT COLOR="#97365B">false</FONT>>, fontname="Courier"];
+    iflinked_assume -> with_linking_assume [label=<  <FONT COLOR="#97365B">true</FONT>>, fontname="Courier"];
+    without_linking_assume -> with_logabsdetjac;
+    with_linking_assume -> with_logabsdetjac;
+    with_logabsdetjac -> return_assume;
+  }
+}
+
diff --git a/docs/src/assets/images/transformations-assume.dot.png b/docs/src/assets/images/transformations-assume.dot.png
diff --git a/docs/src/assets/images/transformations-getindex-with-dist.dot b/docs/src/assets/images/transformations-getindex-with-dist.dot
@@ -0,0 +1,20 @@
+digraph {  
+  # `getindex` block
+  subgraph cluster_getindex {
+    label = "getindex";
+    fontname = "Courier";
+
+    getindex [shape=box, label=< x = getindex(varinfo, <FONT COLOR="#3B6EA8">@varname</FONT>(x), Normal()) >, fontname="Courier"];
+    iflinked_getindex [label=< <FONT COLOR="#3B6EA8">if</FONT> istrans(varinfo, varname) >, fontname="Courier"];
+    without_linking_getindex [shape=box, label="f = from_internal_transform(varinfo, varname, dist)", fontname="Courier"];
+    with_linking_getindex [shape=box, label="f = from_linked_internal_transform(varinfo, varname, dist)", fontname="Courier"];
+    return_getindex [shape=box, label=< <FONT COLOR="#3B6EA8">return</FONT> f(getindex_internal(varinfo, varname)) >, style=dashed, fontname="Courier"];
+
+    getindex -> iflinked_getindex;
+    iflinked_getindex -> without_linking_getindex [label=<  <FONT COLOR="#97365B">false</FONT>>, fontname="Courier"];
+    iflinked_getindex -> with_linking_getindex [label=<  <FONT COLOR="#97365B">true</FONT>>, fontname="Courier"];
+    without_linking_getindex -> return_getindex;
+    with_linking_getindex -> return_getindex;
+  }
+}
+
diff --git a/docs/src/assets/images/transformations-getindex-with-dist.dot.png b/docs/src/assets/images/transformations-getindex-with-dist.dot.png
diff --git a/docs/src/assets/images/transformations-getindex-without-dist.dot b/docs/src/assets/images/transformations-getindex-without-dist.dot
@@ -0,0 +1,20 @@
+digraph {  
+  # `getindex` block
+  subgraph cluster_getindex {
+    label = "getindex";
+    fontname = "Courier";
+
+    getindex [shape=box, label=< x = getindex(varinfo, <FONT COLOR="#3B6EA8">@varname</FONT>(x)) >, fontname="Courier"];
+    iflinked_getindex [label=< <FONT COLOR="#3B6EA8">if</FONT> istrans(varinfo, varname) >, fontname="Courier"];
+    without_linking_getindex [shape=box, label="f = from_internal_transform(varinfo, varname)", fontname="Courier"];
+    with_linking_getindex [shape=box, label="f = from_linked_internal_transform(varinfo, varname)", fontname="Courier"];
+    return_getindex [shape=box, label=< <FONT COLOR="#3B6EA8">return</FONT> f(getindex_internal(varinfo, varname)) >, style=dashed, fontname="Courier"];
+
+    getindex -> iflinked_getindex;
+    iflinked_getindex -> without_linking_getindex [label=<  <FONT COLOR="#97365B">false</FONT>>, fontname="Courier"];
+    iflinked_getindex -> with_linking_getindex [label=<  <FONT COLOR="#97365B">true</FONT>>, fontname="Courier"];
+    without_linking_getindex -> return_getindex;
+    with_linking_getindex -> return_getindex;
+  }
+}
+
diff --git a/docs/src/assets/images/transformations-getindex-without-dist.dot.png b/docs/src/assets/images/transformations-getindex-without-dist.dot.png
diff --git a/docs/src/assets/images/transformations.dot b/docs/src/assets/images/transformations.dot
@@ -0,0 +1,28 @@
+digraph {
+    # Nodes.
+    tilde_node [shape=box, label="x ~ Normal()", fontname="Courier"];
+    base_node [shape=box, label=< vn = <FONT COLOR="#3B6EA8">@varname</FONT>(x)<BR/>dist = Normal()<BR/>x, vi = ... >, fontname="Courier"];
+    assume [shape=box, label="assume(vn, dist, vi)", fontname="Courier"];
+
+    iflinked [label=< <FONT COLOR="#3B6EA8">if</FONT> istrans(vi, vn) >, fontname="Courier"];
+
+    without_linking [shape=box, label="f = from_internal_transform(vi, vn, dist)", styled=dashed, fontname="Courier"];
+    with_linking [shape=box, label="f = from_linked_internal_transform(vi, vn, dist)", styled=dashed, fontname="Courier"];
+
+    with_logabsdetjac [shape=box, label="x, logjac = with_logabsdet_jacobian(f, getindex_internal(vi, vn, dist))", styled=dashed, fontname="Courier"];
+    return [shape=box, label=< <FONT COLOR="#3B6EA8">return</FONT> x, logpdf(dist, x) - logjac, vi >, styled=dashed, fontname="Courier"];
+
+    # Edges.
+    tilde_node -> base_node [style=dashed, label=<  <FONT COLOR="#3B6EA8">@model</FONT>>, fontname="Courier"]
+    base_node -> assume [style=dashed, label="  tilde-pipeline", fontname="Courier"];
+
+    assume -> iflinked;
+
+    iflinked -> without_linking [label=<  <FONT COLOR="#97365B">false</FONT>>, fontname="Courier"];
+    iflinked -> with_linking [label=<  <FONT COLOR="#97365B">true</FONT>>, fontname="Courier"];
+
+    without_linking -> with_logabsdetjac;
+    with_linking -> with_logabsdetjac;
+
+    with_logabsdetjac -> return;
+}
diff --git a/docs/src/assets/images/transformations.dot.png b/docs/src/assets/images/transformations.dot.png