From 9b057c011e33eca6b4205178e2ac29ff6d5461fb Mon Sep 17 00:00:00 2001
From: Chad Scherrer <chad.scherrer@gmail.com>
Date: Tue, 24 May 2022 12:58:22 -0700
Subject: [PATCH] Dev (#332)

* update Project.toml

* s/logpdf/logdensity/

* update deps

* scratchpad

* version bump

* representative => rootmeasure

* update dependency versions

* reduce dependencies

* update deps

* scratchpad

* representative => rootmeasure

* update dependency versions

* reduce dependencies

* `as` methods for `xform`

* cleanup

* require latest MeasureTheory

* dorp old distributions code

* drop old iid code

* drop extra space

* limit deps to three newest releases

* update dynamichmc

* add Aqua

* bump version

* Better `predict` method

* withmeasures(::ConditionalModel)

* update dependencies

* updating symbolics

* some updates to symbolics

* hmm example

* update MeasureBase bound

* better dispatch for `predict`

* drop redundant method

* remove whitespace

* bump version

* update for upcoming MeasureTheory release

* small change to toposort

* have `iid` use powermeasure

* tests passing!

* iid(n::Integer...)

* update dependencies

* start logdensity with partialstatic

* update test to account for partialstatic

* add `insupport`

* clean up insupport

* speed up model building

* updates

* starting updates for new MeasureTheory stuff

* get tests passing

* Kleisli => TransitionKernel

* bump MB dependency version

* Kleisli => TransitionKernel

* tests passing

* update dependencies

* update CI

* Should only need BayesianLinearRegression for testing

* bump version
---
 .github/workflows/ci.yml                    |  3 +-
 Project.toml                                | 38 ++++++------
 demos/boundmodels.jl                        |  2 +-
 demos/hmm.jl                                |  2 +-
 demos/repeated-measurements.jmd             |  6 +-
 demos/repeated-measurements.md              |  6 +-
 docs/src/misc.md                            |  2 +-
 scratchpad/MCM.jl                           |  4 +-
 scratchpad/SampleChainsDynamicHMC-demo.jl   |  4 +-
 scratchpad/advancedhmc.jl                   |  2 +-
 scratchpad/bouncy.jl                        |  4 +-
 scratchpad/codegen.jl                       |  2 +-
 scratchpad/composition.jl                   |  8 +--
 scratchpad/composition2.jl                  |  2 +-
 scratchpad/example-linear-regression.jl     |  2 +-
 scratchpad/flux.jl                          |  2 +-
 scratchpad/gps.jl                           |  2 +-
 scratchpad/hmm.jl                           | 44 ++++++++++----
 scratchpad/repeated-measurements-example.jl | 10 ++--
 scratchpad/s4.jl                            | 10 ++--
 scratchpad/stheno.jl                        |  2 +-
 scratchpad/trying_out_tuplevectors.jl       |  4 +-
 scratchpad/vis.jl                           | 10 ++--
 scratchpad/vis2.jl                          | 22 +++----
 src/Soss.jl                                 | 37 +++++++-----
 src/bayeslm.jl                              |  4 +-
 src/core/models/abstractmodel.jl            | 13 +++--
 src/core/models/conditional.jl              |  2 +-
 src/core/toposort.jl                        |  4 +-
 src/core/utils.jl                           | 39 +++++++++----
 src/distributions/dist.jl                   |  4 +-
 src/distributions/iid.jl                    |  6 +-
 src/exponentialfamilies.jl                  |  6 +-
 src/importance.jl                           | 10 ++--
 src/inference/advancedhmc.jl                |  4 +-
 src/inference/dynamicHMC.jl                 | 12 ++--
 src/inference/power-posterior.jl            |  2 +-
 src/inference/rejection.jl                  | 20 +++----
 src/optim.jl                                |  2 +-
 src/primitives/{xform.jl => as.jl}          | 36 ++++++------
 src/primitives/basemeasure.jl               | 10 ++--
 src/primitives/bijectors.jl                 |  2 +-
 src/primitives/insupport.jl                 | 49 ++++++++++++++++
 src/primitives/likelihood-weighting.jl      |  2 +-
 src/primitives/logdensity.jl                | 65 +++------------------
 src/samplechains/dynamichmc.jl              |  4 +-
 src/scratch/bouncy.jl                       |  4 +-
 src/scratch/zigzag.jl                       |  2 +-
 src/symbolic/symbolic.jl                    |  4 +-
 src/transforms/utils.jl                     |  3 +-
 test/linear-model.jl                        |  4 +-
 test/runtests.jl                            | 28 ++++++---
 52 files changed, 318 insertions(+), 253 deletions(-)
 rename src/primitives/{xform.jl => as.jl} (64%)
 create mode 100644 src/primitives/insupport.jl

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 18f68d48..280cb579 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -17,7 +17,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.5'
+          - '1.6'
+          - '1.7'
           - '1'
         os:
           - ubuntu-latest
diff --git a/Project.toml b/Project.toml
index 44b8f5c0..4b181067 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,15 +1,16 @@
 name = "Soss"
 uuid = "8ce77f84-9b61-11e8-39ff-d17a774bf41c"
 author = ["Chad Scherrer <chad.scherrer@gmail.com>"]
-version = "0.20.9"
+version = "0.21.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb"
-IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"
+IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 JuliaVariables = "b14d175d-62b4-44ba-8fb7-3064adc8c3ec"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
@@ -30,45 +31,46 @@ SimpleGraphs = "55797a34-41de-5266-9ec1-32ac4eb504d3"
 SimplePartitions = "ec83eff0-a5b5-5643-ae32-5cbf6eedec9d"
 SimplePosets = "b2aef97b-4721-5af9-b440-0bad754dc5ba"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-SymbolicCodegen = "fc9b0551-4f1b-42e1-8440-e8a535f5a551"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"
 TupleVectors = "615932cf-77b6-4358-adcd-5b7eba981d7e"
 
 [compat]
-ArrayInterface = "3"
+ArrayInterface = "5,6"
+DensityInterface = "0.4"
 DiffResults = "1"
 Distributions = "0.23, 0.24, 0.25"
-FillArrays = "0.10, 0.11, 0.12"
+FillArrays = "0.11, 0.12, 0.13"
 GeneralizedGenerated = "0.3"
-IRTools = "0.4"
+IfElse = "0.1"
 JuliaVariables = "0.2"
 MLStyle = "0.3,0.4"
 MacroTools = "0.5"
 MappedArrays = "0.3, 0.4"
-MeasureBase = "0.5"
-MeasureTheory = "0.13"
-NamedTupleTools = "0.12, 0.13"
-NestedTuples = "0.3"
-RecipesBase = "0.7,0.8, 1"
+MeasureBase = "0.9"
+MeasureTheory = "0.16"
+NamedTupleTools = "0.12, 0.13, 0.14"
+NestedTuples = "0.3.9"
+RecipesBase = "1"
 Reexport = "1"
 Requires = "1"
 RuntimeGeneratedFunctions = "0.5"
 SampleChains = "0.5"
 SimpleGraphs = "0.5, 0.6, 0.7"
 SimplePartitions = "0.2, 0.3"
-SimplePosets = "0.0, 0.1"
-SpecialFunctions = "0.9, 0.10, 1"
+SimplePosets = "0.1"
+SpecialFunctions = "1, 2"
+Static = "0.5, 0.6"
 StatsBase = "0.33"
-StatsFuns = "0.9"
-SymbolicCodegen = "0.2"
-SymbolicUtils = "0.15, 0.16, 0.17"
-TransformVariables = "0.4"
+StatsFuns = "0.9, 1"
+SymbolicUtils = "0.17, 0.18, 0.19"
+TransformVariables = "0.5, 0.6"
 TupleVectors = "0.1"
-julia = "1.5"
+julia = "1.6"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
diff --git a/demos/boundmodels.jl b/demos/boundmodels.jl
index 4e420561..c7b02c08 100644
--- a/demos/boundmodels.jl
+++ b/demos/boundmodels.jl
@@ -10,7 +10,7 @@ m(μ=1.0)
 
 rand(m(μ=1.0))
 
-logdensity(m(μ=1.0),(x=0.4,))
+logdensityof(m(μ=1.0),(x=0.4,))
 # rand(m1)
 
 weightedSample(m(μ=1.0),(x=3,))
diff --git a/demos/hmm.jl b/demos/hmm.jl
index 8f29fb44..3d54cf04 100644
--- a/demos/hmm.jl
+++ b/demos/hmm.jl
@@ -53,7 +53,7 @@ end
     end
 end
 
-# function Distributions.logdensity(c::Chain, xs)
+# function Distributions.logdensityof(c::Chain, xs)
 
 c = Chain(s0, step)
 
diff --git a/demos/repeated-measurements.jmd b/demos/repeated-measurements.jmd
index 872c2113..bf692ea8 100644
--- a/demos/repeated-measurements.jmd
+++ b/demos/repeated-measurements.jmd
@@ -92,8 +92,8 @@ for t in symlogdensity(m()).evalf(3).args
 
 <!--
 using BenchmarkTools
-@btime logdensity(m(),truth)
-@btime logdensity(m(),truth, codegen)
+@btime logdensityof(m(),truth)
+@btime logdensityof(m(),truth, codegen)
 
 f1 = Soss._codegen(m, true);
 f2 = Soss._codegen(m,false);
@@ -105,7 +105,7 @@ codegen(m(),truth)
 
 
 
-logdensity(m(), merge(truth, (p_bad=shuffle(truth.p_bad),)), codegen)
+logdensityof(m(), merge(truth, (p_bad=shuffle(truth.p_bad),)), codegen)
 
 
 @time result = dynamicHMC(m(), (y=truth.y,), codegen) ;
diff --git a/demos/repeated-measurements.md b/demos/repeated-measurements.md
index e547533e..8966b627 100644
--- a/demos/repeated-measurements.md
+++ b/demos/repeated-measurements.md
@@ -192,8 +192,8 @@ _j1 = 1 _j2 = 1
 
 <!--
 using BenchmarkTools
-@btime logdensity(m(),truth)
-@btime logdensity(m(),truth, codegen)
+@btime logdensityof(m(),truth)
+@btime logdensityof(m(),truth, codegen)
 
 f1 = Soss._codegen(m, true);
 f2 = Soss._codegen(m,false);
@@ -205,7 +205,7 @@ codegen(m(),truth)
 
 
 
-logdensity(m(), merge(truth, (p_bad=shuffle(truth.p_bad),)), codegen)
+logdensityof(m(), merge(truth, (p_bad=shuffle(truth.p_bad),)), codegen)
 
 
 @time result = dynamicHMC(m(), (y=truth.y,), codegen) ;
diff --git a/docs/src/misc.md b/docs/src/misc.md
index 82602a85..04109dc4 100644
--- a/docs/src/misc.md
+++ b/docs/src/misc.md
@@ -65,7 +65,7 @@ function sourceWeightedSample(_data)
         proc(_m, st :: LineNumber) = nothing
 
         function proc(_m, st :: Sample)
-            st.x ∈ _datakeys && return :(_ℓ += logdensity($(st.rhs), $(st.x)))
+            st.x ∈ _datakeys && return :(_ℓ += logdensity_def($(st.rhs), $(st.x)))
             return :($(st.x) = rand($(st.rhs)))
         end
 
diff --git a/scratchpad/MCM.jl b/scratchpad/MCM.jl
index 15221a53..f28f1c4b 100644
--- a/scratchpad/MCM.jl
+++ b/scratchpad/MCM.jl
@@ -7,7 +7,7 @@ function expect(x,ℓ)
 end
 
 x = Particles(5000,TDist(3))
-ℓ = logpdf(Normal(),x) - logpdf(TDist(3), x)
+ℓ = logdensityof(Normal(),x) - logdensityof(TDist(3), x)
 
 expect(x,ℓ)
 expect(x^2,ℓ)
@@ -26,7 +26,7 @@ expect(z^4,ℓ)
 # And now a new weighted variable
 
 y = Particles(5000,TDist(3))
-ℓ += logpdf(Normal(),y) - logpdf(TDist(3), y)
+ℓ += logdensityof(Normal(),y) - logdensityof(TDist(3), y)
 
 expect(y,ℓ)
 expect(y^2,ℓ)
diff --git a/scratchpad/SampleChainsDynamicHMC-demo.jl b/scratchpad/SampleChainsDynamicHMC-demo.jl
index 603f081c..73f48c28 100644
--- a/scratchpad/SampleChainsDynamicHMC-demo.jl
+++ b/scratchpad/SampleChainsDynamicHMC-demo.jl
@@ -44,8 +44,8 @@ plot(p, exp.([ℓ((p=pj,)) for pj in p]))
 
 using TransformVariables
 
-ℓ(x) = logdensity(Beta(4,3), x.p)
-t = xform(post)
+ℓ(x) = logdensity_def(Beta(4,3), x.p)
+t = as(post)
 chain = initialize!(DynamicHMCChain, ℓ, t)
 drawsamples!(chain, 10000)
 plot(ash(chain.p))
diff --git a/scratchpad/advancedhmc.jl b/scratchpad/advancedhmc.jl
index 164bc079..2569f4ad 100644
--- a/scratchpad/advancedhmc.jl
+++ b/scratchpad/advancedhmc.jl
@@ -6,7 +6,7 @@ D = 10; initial_θ = rand(D)
 dist = Normal() ^ D
 
 # Define the target distribution
-ℓπ(θ) = logdensity(dist, θ)
+ℓπ(θ) = logdensity_def(dist, θ)
 
 # Set the number of samples to draw and warmup iterations
 n_samples, n_adapts = 2_000, 1_000
diff --git a/scratchpad/bouncy.jl b/scratchpad/bouncy.jl
index 12330479..c90563ae 100644
--- a/scratchpad/bouncy.jl
+++ b/scratchpad/bouncy.jl
@@ -17,9 +17,9 @@ Returns a `Trace` object.
 """
 function bouncy(m::ConditionalModel, T = 1000.0;  c=10.0, λref=0.1, ρ=0.0, adapt=false) where {A,B}
 
-    ℓ(pars) = logdensity(m, pars)
+    ℓ(pars) = logdensity_def(m, pars)
 
-    t = xform(m)
+    t = as(m)
 
     function f(x)
         (θ, logjac) = transform_and_logjac(t, x)
diff --git a/scratchpad/codegen.jl b/scratchpad/codegen.jl
index 778a33b6..31c04d69 100644
--- a/scratchpad/codegen.jl
+++ b/scratchpad/codegen.jl
@@ -8,7 +8,7 @@ m = @model σ begin
 
 x = rand(m(σ=3.0));
 
-sourceSymlogdensity(m(σ=3.0) | (;x))
+sourceSymlogdensityof(m(σ=3.0) | (;x))
 
 s = symlogdensity(m(σ=3.0))
 
diff --git a/scratchpad/composition.jl b/scratchpad/composition.jl
index 3a0ec70f..1be4d8d8 100644
--- a/scratchpad/composition.jl
+++ b/scratchpad/composition.jl
@@ -33,17 +33,17 @@ dynamicHMC(post)
 
 
 
-logdensity(post, truth)
+logdensity_def(post, truth)
 
 speed = truth.speed
 x = truth.x
 v = truth.v
 
 _ℓ = 0.0
-_ℓ += logdensity(Normal(0.0, 100.0), speed)
+_ℓ += logdensity_def(Normal(0.0, 100.0), speed)
 speed = predict(Normal(0.0, 100.0), speed)
-_ℓ += logdensity(gps(; speed, n, t), x)
+_ℓ += logdensity_def(gps(; speed, n, t), x)
 x = predict(gps(; speed, n, t), x)
-_ℓ += logdensity(radar(; speed, n), v)
+_ℓ += logdensity_def(radar(; speed, n), v)
 v = predict(radar(; speed, n), v)
 _ℓ
diff --git a/scratchpad/composition2.jl b/scratchpad/composition2.jl
index dddbcfe1..514dc5ce 100644
--- a/scratchpad/composition2.jl
+++ b/scratchpad/composition2.jl
@@ -29,7 +29,7 @@ rand(m())
 
 x = randn(2)
 
-xform(m() | (x=x, μ = (z = 1.0,)))
+as(m() | (x=x, μ = (z = 1.0,)))
 
 dynamicHMC(m() | (x=x, μ = (z = 1.0,)))
 
diff --git a/scratchpad/example-linear-regression.jl b/scratchpad/example-linear-regression.jl
index cd35d02c..d73b12ba 100644
--- a/scratchpad/example-linear-regression.jl
+++ b/scratchpad/example-linear-regression.jl
@@ -130,7 +130,7 @@ y_true - particles(y_ppc)
 
  # `rand`  turns each `v ~ dist` into `v = rand(dist)`, finally outputting the `NamedTuple` of all values it has seen.
 
-# `logdensity` steps through the same program, but instead accumulates a log-density. It begins by initializing `_ℓ = 0.0`. Then at each step, it turns `v ~ dist` into `_ℓ += logdensity(dist, v)`, before finally returning `_ℓ`.
+# `logdensity` steps through the same program, but instead accumulates a log-density. It begins by initializing `_ℓ = 0.0`. Then at each step, it turns `v ~ dist` into `_ℓ += logdensity_def(dist, v)`, before finally returning `_ℓ`.
 
 # Note that I said "turns into" instead of "interprets". Soss uses [`GG.jl`](https://github.com/thautwarm/GG.jl) to generate specialized code for a given model, inference primitive (like `rand` and `logdensity`), and type of data.
 
diff --git a/scratchpad/flux.jl b/scratchpad/flux.jl
index 0f5ade5f..b0f4d81b 100644
--- a/scratchpad/flux.jl
+++ b/scratchpad/flux.jl
@@ -5,7 +5,7 @@
 #     t = getTransform(model)
 #     z = randn(t.dimension) |> Flux.param
 
-#     fpre = @eval $(logdensity(model))
+#     fpre = @eval $(logdensity_def(model))
 #     f(par, data) = Base.invokelatest(fpre, par, data)
 
 #     loss(data) = -f(transform(t, z), data)
diff --git a/scratchpad/gps.jl b/scratchpad/gps.jl
index 6b94b7b4..32b66432 100644
--- a/scratchpad/gps.jl
+++ b/scratchpad/gps.jl
@@ -8,7 +8,7 @@ y = sinpi.(x) + 0.3*randn(100)
 
 import TransformVariables
 const TV = TransformVariables 
-Soss.xform(gp::AbstractGPs.FiniteGP, _data=NamedTuple()) = TV.as(Array, TV.asℝ,size(gp)...)
+Soss.as(gp::AbstractGPs.FiniteGP, _data=NamedTuple()) = as(Array, asℝ,size(gp)...)
 
 sqexpkernel(alpha::Real, rho::Real) = 
 alpha^2 * transform(SEKernel(), 1/(rho*sqrt(2)))
diff --git a/scratchpad/hmm.jl b/scratchpad/hmm.jl
index ee72fb9a..18b7e764 100644
--- a/scratchpad/hmm.jl
+++ b/scratchpad/hmm.jl
@@ -1,16 +1,41 @@
 using MeasureTheory
 using Base.Iterators
 using Statistics
-
+using Soss
 using Random
-rng = Random.Xoshiro(3)
 
-x = Chain(Normal()) do xj Normal(μ=xj) end
-xobs = rand(rng, x)
-y = For(xobs) do xj Poisson(logλ=xj) end
-yobs = rand(rng, y)
-xv = take(xobs, 10) |> collect
-yv = take(yobs, 10) |> collect
+# 5 7 10
+
+rng = Random.Xoshiro(12)
+
+m = @model begin
+    latent ~ Chain(Normal(μ=1)) do x Normal(μ=x, σ=0.2) end
+    observed ~ For(latent) do x Poisson(logλ=x) end
+end
+
+truth = rand(rng, m())
+
+xv = take(truth.latent, 100) |> collect
+yv = take(truth.observed, 100) |> collect
+
+
+using Plots
+
+plt = scatter(yv, label="observations")
+plot!(exp.(xv), lw=3, label="latent process")
+
+using Statistics
+
+scatter(exp.(xv), poiscdf.(exp.(xv), yv), label=false)
+
+x = Chain(Normal(μ=1)) do xj Normal(μ=xj, σ=0.2) end
+
+
+xvals = rand(rng, x)
+y = For(xobs) do xj Poisson(logλ= xj) end
+yvals = rand(rng, y)
+xv = take(xobs, 100) |> collect
+yv = take(yobs, 100) |> collect
 
 take(xobs.parent, 10) |> collect
 take(yobs.parent, 10) |> collect
@@ -19,7 +44,6 @@ take(yobs.parent, 10) |> collect
 exp.(xv)
 yv
 
-# using Plots
 
 # plt = scatter(normcdf.(xv, 1, yv), label=false)
 
@@ -44,5 +68,5 @@ truth = rand(rng, m())
 xobs = take(truth.x, 10) |> collect
 yobs = take(truth.y, 10) |> collect
 
-logdensity(m(), (x=xobs, y=yobs))
+logdensity_def(m(), (x=xobs, y=yobs))
 
diff --git a/scratchpad/repeated-measurements-example.jl b/scratchpad/repeated-measurements-example.jl
index 209e8b63..5b217084 100644
--- a/scratchpad/repeated-measurements-example.jl
+++ b/scratchpad/repeated-measurements-example.jl
@@ -77,7 +77,7 @@ function elbo(x)
 
     qλ = q(r=5,s=14, λα=λα)
     r = rand(qλ)
-    logdensity(posterior, r) - logdensity(qλ, r)
+    logdensity_def(posterior, r) - logdensity_def(qλ, r)
 end
 
 elbo(x)
@@ -139,7 +139,7 @@ x0 = zeros(48)
 
 
 
-logdensity(p(r=5, s=14) | (;y), rand(q(r=5,s=14, λ=λ.λ)))
+logdensity_def(p(r=5, s=14) | (;y), rand(q(r=5,s=14, λ=λ.λ)))
 
 
 
@@ -202,7 +202,7 @@ logdensity(p(r=5, s=14) | (;y), rand(q(r=5,s=14, λ=λ.λ)))
 
 function elbo(p, args, obs, q, λ)
     qargs = merge(args, λ)
-    mean(logpdf(p(args), merge(particles(q(qargs)), obs))) + entropy(q(args))
+    mean(logdensityof(p(args), merge(particles(q(qargs)), obs))) + entropy(q(args))
 end
 
 
@@ -239,7 +239,7 @@ using AdvancedVI
 
 bad = rand(Bernoulli(0.5) |> iid(14))
 
-tr = xform(p(r=5,s=14), (y=y,));
+tr = as(p(r=5,s=14), (y=y,));
 d = tr.dimension
 
 getq(λ) = MvNormal(λ[1:d], exp.(λ[d .+ (1:d)]))
@@ -249,7 +249,7 @@ advi = ADVI(10, 10_000)
 
 
 function logπ(θ)
-    logpdf(p(r=5,s=14),  merge((y=y,), tr(θ)))
+    logdensityof(p(r=5,s=14),  merge((y=y,), tr(θ)))
 end
 
 logπ(randn(24))
diff --git a/scratchpad/s4.jl b/scratchpad/s4.jl
index 33631d3d..50802abf 100644
--- a/scratchpad/s4.jl
+++ b/scratchpad/s4.jl
@@ -18,11 +18,11 @@ end
 
 
 
-function logdensity(mc::Chain, x)
+function logdensity_def(mc::Chain, x)
     μ = mc.μ
     ℓ = 0.0
     for xj in x
-        ℓ += logdensity(μ, xj)
+        ℓ += logdensity_def(μ, xj)
         μ = mc.κ(xj)
     end
     return ℓ
@@ -30,7 +30,7 @@ end
 
 mc = Chain(Normal()) do x Normal(μ=x) end
 
-logdensity(mc, randn(100))
+logdensity_def(mc, randn(100))
 
 
 using Soss
@@ -82,7 +82,7 @@ rand(d2)
 
 simulate(d)
 
-t = xform(d);
+t = as(d);
 t(randn(4))
 
 
@@ -101,7 +101,7 @@ t(randn(4))
 # julia> rand(d)
 # -3.0414465047589037
 
-# julia> t = xform(d)
+# julia> t = as(d)
 # TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.Identity, TransformVariables.Identity}}}, TransformVariables.Identity}}}, TransformVariables.Identity}}}((x = TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.Identity, TransformVariables.Identity}}}, TransformVariables.Identity}}}((x = TransformVariables.TransformTuple{NamedTuple{(:x, :y), Tuple{TransformVariables.Identity, TransformVariables.Identity}}}((x = asℝ, y = asℝ), 2), y = asℝ), 3), y = asℝ), 4)
 
 # julia> t(randn(4))
diff --git a/scratchpad/stheno.jl b/scratchpad/stheno.jl
index 589ae741..5937e51f 100644
--- a/scratchpad/stheno.jl
+++ b/scratchpad/stheno.jl
@@ -23,7 +23,7 @@ x = [1:50;];
 
 rand(m(x=x))
 
-# logpdf(m(x=x), (;α=9.0, β=0.1, y=births2.births))
+# logdensityof(m(x=x), (;α=9.0, β=0.1, y=births2.births))
 
 post = dynamicHMC(m(x=x), (;y=births2.births[1:50]))
 
diff --git a/scratchpad/trying_out_tuplevectors.jl b/scratchpad/trying_out_tuplevectors.jl
index 196765d2..ac2a18bb 100644
--- a/scratchpad/trying_out_tuplevectors.jl
+++ b/scratchpad/trying_out_tuplevectors.jl
@@ -39,9 +39,9 @@ y = rand(p(k=k,x=x))
 
 post = m(k=k, x=x) | (;y)
 
-ℓ(nt) = logdensity(post, nt) ;
+ℓ(nt) = logdensity_def(post, nt) ;
 
-t = xform(post);
+t = as(post);
 
 chains = initialize!(DynamicHMCChain, ℓ, t);
 drawsamples!(chains, 199)
diff --git a/scratchpad/vis.jl b/scratchpad/vis.jl
index 06756ac9..a5a0b91d 100644
--- a/scratchpad/vis.jl
+++ b/scratchpad/vis.jl
@@ -84,7 +84,7 @@ end
 
 @inline function visStep(N,logp,q)
     x = Particles(N,q)
-    ℓ = logp(x) - logdensity(q,x)
+    ℓ = logp(x) - logdensity_def(q,x)
     μ = expect(x,ℓ)
     Σ = cholesky(Positive, expect(0.5 * self_outer(x-μ),ℓ)) |> PDMat
     return (x,ℓ,μ,Σ)
@@ -93,7 +93,7 @@ end
 function runInference(m; kwargs...)
     ℓp_pre = sourceLogdensity(m) |> eval
     ℓp(θ) = Base.invokelatest(ℓp_pre, θ) 
-    t = xform(m; kwargs...)
+    t = as(m; kwargs...)
     d = t.dimension
     tnames = propertynames(t.transformations)
 
@@ -102,7 +102,7 @@ function runInference(m; kwargs...)
 
     function logp(x::Vector{T} where T)
         prep(θ) = merge(kwargs, θ)
-        f = transform_logdensity(t, ℓp ∘ prep, x)
+        f = transform_logdensity_def(t, ℓp ∘ prep, x)
     end
 
     logp(x) = logp([x])
@@ -119,7 +119,7 @@ function runInference(m; kwargs...)
     N = 1000
     q = fit_mle(MvNormal, asmatrix(x))
     x = Particles(N,q)
-    ℓ = logp(x) - logdensity(q,x)
+    ℓ = logp(x) - logdensity_def(q,x)
 
 
     plts = []
@@ -142,7 +142,7 @@ function runInference(m; kwargs...)
 
     x = Particles(N,q)
     θ = t(x)
-    ℓ = logp(x) - logdensity(q,x)
+    ℓ = logp(x) - logdensity_def(q,x)
     (θ,q,ℓ,elbo,neff)
 end
 
diff --git a/scratchpad/vis2.jl b/scratchpad/vis2.jl
index 9448a9e3..963b77b6 100644
--- a/scratchpad/vis2.jl
+++ b/scratchpad/vis2.jl
@@ -20,10 +20,10 @@ end
 function fromObs(x,y) 
     function logp(α,β)
         ℓ = 0.0
-        ℓ += logpdf(Normal(0,1), α)
-        ℓ += logpdf(Normal(0,2), β)
+        ℓ += logdensityof(Normal(0,1), α)
+        ℓ += logdensityof(Normal(0,2), β)
         yhat = α .+ β .* x
-        ℓ += sum(logpdf.(Normal.(yhat, 1), y) )
+        ℓ += sum(logdensityof.(Normal.(yhat, 1), y) )
         ℓ
     end
 end
@@ -63,14 +63,14 @@ function runInference(x,y,logp)
     q = MvNormal(2,100000.0) # Really this would be fit from a sample from the prior
     α,β = Particles(N,q)
     m = asmatrix(α,β)
-    ℓ = sum(logp(α,β)) - Particles(logpdf(q,m))
+    ℓ = sum(logp(α,β)) - Particles(logdensityof(q,m))
 
     numiters = 60
     elbo = Vector{Float64}(undef, numiters)
     for j in 1:numiters
         α,β = Particles(N,q)
         m = asmatrix(α,β)
-        ℓ = logp(α,β) - Particles(logpdf(q,m))
+        ℓ = logp(α,β) - Particles(logdensityof(q,m))
         elbo[j] = mean(ℓ) 
         ss = suffstats(MvNormal, m,  exp(ℓ - maximum(ℓ)).particles .+ 1/N)
         q = fit_mle(MvNormal, ss)
@@ -129,7 +129,7 @@ q = MvNormal(2,100.0)
 
 
 function initialize(jointdist, obs)
-    tr = xform(jointdist, obs)
+    tr = as(jointdist, obs)
     pars = keys(tr.transformations)
     m = jointdist.model
 
@@ -140,7 +140,7 @@ function initialize(jointdist, obs)
     q = inverse.(tr, samples) 
 
 
-    ℓ = [logpdf(pred(merge(args, s)), obs) for s in samples]
+    ℓ = [logdensityof(pred(merge(args, s)), obs) for s in samples]
 end
 
 
@@ -166,14 +166,14 @@ function updateSample!(s,q)
 end
 
 function updateWeights!(w, s, p, tr)
-    map!(θ -> logpdf(p, θ), w, tr.(eachcol(s)))
+    map!(θ -> logdensityof(p, θ), w, tr.(eachcol(s)))
 end
 
 inverse(tr, select(rand(m(x=x)), pars))
 
 
-Soss.prior(m,:y) |> xform
+Soss.prior(m,:y) |> as
 
-Soss.sourcexform(Soss.prior(m,:y))
+Soss.sourceas(Soss.prior(m,:y))
 
-xform()
+as()
diff --git a/src/Soss.jl b/src/Soss.jl
index daaab8d8..403bc809 100644
--- a/src/Soss.jl
+++ b/src/Soss.jl
@@ -5,21 +5,21 @@ using Random
 using Reexport: @reexport
 
 @reexport using StatsFuns
-using MeasureTheory
+@reexport using MeasureTheory
 using MeasureBase: productmeasure, Returns
 
-export For
+import DensityInterface: logdensityof
+import DensityInterface: densityof
+import DensityInterface: DensityKind
+using DensityInterface
 
 using NamedTupleTools
 using SampleChains
-using SymbolicCodegen
+# using SymbolicCodegen
 
 using SymbolicUtils: Symbolic
 const MaybeSym{T} = Union{T, Symbolic{T}}
 
-# MeasureTheory.For(f, dims::MaybeSym{<: Integer}...) = ProductMeasure(mappedarray(i -> f(Tuple(i)...), CartesianIndices(dims))) 
-# MeasureTheory.For(f, dims::MaybeSym{<: AbstractUnitRange}...) = ProductMeasure(mappedarray(i -> f(Tuple(i)...), CartesianIndices(dims))) 
-
 import MacroTools: prewalk, postwalk, @q, striplines, replace, @capture
 import MacroTools
 import MLStyle
@@ -28,9 +28,10 @@ import MLStyle
 
 using Requires
 using ArrayInterface: StaticInt
+using Static
 
-
-using TransformVariables: as, asℝ, as𝕀, asℝ₊
+using IfElse: ifelse
+using TransformVariables: asℝ, as𝕀, asℝ₊
 import TransformVariables
 const TV = TransformVariables
 
@@ -38,9 +39,11 @@ using SimplePosets: SimplePoset
 
 using RuntimeGeneratedFunctions
 RuntimeGeneratedFunctions.init(@__MODULE__)
+using MeasureBase: AbstractTransitionKernel
 
 using MeasureTheory: ∞
-
+@reexport using MeasureTheory
+import MeasureTheory: as
 
 """
 we use this to avoid introduce static type parameters
@@ -48,6 +51,9 @@ for generated functions
 """
 _unwrap_type(a::Type{<:Type}) = a.parameters[1]
 
+import GeneralizedGenerated as GG
+
+
 include("noted.jl")
 include("core/models/abstractmodel.jl")
 include("core/statement.jl")
@@ -70,8 +76,9 @@ include("distributions/iid.jl")
 include("primitives/rand.jl")
 include("primitives/simulate.jl")
 include("primitives/logdensity.jl")
-include("primitives/xform.jl")
+include("primitives/as.jl")
 include("primitives/likelihood-weighting.jl")
+include("primitives/insupport.jl")
 # include("primitives/gg.jl")
 # @init @require Bijectors="76274a88-744f-5084-9051-94815aaf08c4" begin
 #     include("primitives/bijectors.jl")
@@ -88,11 +95,11 @@ include("transforms/utils.jl")
 include("transforms/basictransforms.jl")
 include("transforms/withmeasures.jl")
 
-include("symbolic/symcall.jl")
-include("symbolic/symify.jl")
-include("symbolic/rules.jl")
-include("symbolic/symbolic.jl")
-include("symbolic/codegen.jl")
+# include("symbolic/symcall.jl")
+# include("symbolic/symify.jl")
+# include("symbolic/rules.jl")
+# include("symbolic/symbolic.jl")
+# include("symbolic/codegen.jl")
 
 # include("particles.jl")
 include("plots.jl")
diff --git a/src/bayeslm.jl b/src/bayeslm.jl
index 74fd0869..e2c6e7a2 100644
--- a/src/bayeslm.jl
+++ b/src/bayeslm.jl
@@ -28,7 +28,7 @@ end
 function makeℓ(X, y, pr)
     k = size(X,2)
     ll = loglik(X,y)
-    ℓ(pars) = ll(pars.α, pars.β, pars.σ) + logdensity(pr(k=k), pars)
+    ℓ(pars) = ll(pars.α, pars.β, pars.σ) + logdensity_def(pr(k=k), pars)
 end
 
 function bayeslm(
@@ -41,7 +41,7 @@ function bayeslm(
     kwargs...,
 )
     ℓ = makeℓ(X, y, pr)
-    t = xform(pr(k=size(X,2)))
+    t = as(pr(k=size(X,2)))
 
 
     chain = initialize!(DynamicHMCChain, ℓ, t)
diff --git a/src/core/models/abstractmodel.jl b/src/core/models/abstractmodel.jl
index 1f47ef9b..bf7b3fbc 100644
--- a/src/core/models/abstractmodel.jl
+++ b/src/core/models/abstractmodel.jl
@@ -15,13 +15,13 @@ N gives the Names of arguments (each a Symbol)
 B gives the Body, as an Expr
 M gives the Module where the model is defined
 """
-abstract type AbstractModel{A,B,M,Args,Obs} <: AbstractMeasure end
+abstract type AbstractModel{A,B,M,Args,Obs} <: AbstractTransitionKernel end
 
-argstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = AT
-argstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = AT
+argstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = A
+argstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = A
 
-bodytype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = BT
-bodytype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = BT
+bodytype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = B
+bodytype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = B
 
 getmodule(::Type{AbstractModel{A,B,M,Args,Obs}}) where  {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = from_type(M)
 getmodule(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = from_type(M)
@@ -34,3 +34,6 @@ argvalstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}}
 
 obstype(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = Obs
 obstype(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = Obs
+
+body(::AbstractModel{A,B,M,Args,Obs}) where {A,B,M,Args,Obs} = from_type(B)
+body(::Type{AM}) where {A,B,M,Args,Obs,AM<:AbstractModel{A,B,M,Args,Obs}} = from_type(B)
\ No newline at end of file
diff --git a/src/core/models/conditional.jl b/src/core/models/conditional.jl
index 89bdca89..caccc1e1 100644
--- a/src/core/models/conditional.jl
+++ b/src/core/models/conditional.jl
@@ -35,7 +35,7 @@ ConditionalModel(m::Model) = ConditionalModel(m,NamedTuple(), NamedTuple())
 
 (m::Model)(;argvals...)= m((;argvals...))
 
-(m::Model)(args...) = m(NamedTuple{Tuple(m.args)}(args))
+(m::Model)(args...) = m(argstype(m)(args))
 
 import Base
 
diff --git a/src/core/toposort.jl b/src/core/toposort.jl
index 7704d32a..58e25e55 100644
--- a/src/core/toposort.jl
+++ b/src/core/toposort.jl
@@ -12,14 +12,14 @@ end
 function toposort(g::SimpleDigraph{T}) where {T}
     g = deepcopy(g)
     order = T[]
-    s = collect(filter(v -> SimpleGraphs.in_deg(g, v) == 0, vlist(g)))
+    s = Symbol[v for v in vlist(g) if SimpleGraphs.in_deg(g, v) == 0]
 
     while !isempty(s)
         u = pop!(s)
         push!(order, u)
 
         for v in out_neighbors(g, u)
-            delete!(g, u, v)
+            SimpleGraphs.delete!(g, u, v)
             if SimpleGraphs.in_deg(g, v) == 0
                 push!(s, v)
             end
diff --git a/src/core/utils.jl b/src/core/utils.jl
index 434611d6..25d3ee7e 100644
--- a/src/core/utils.jl
+++ b/src/core/utils.jl
@@ -202,27 +202,42 @@ function loadvals(argstype, datatype)
 end
 
 function loadvals(argstype, datatype, parstype)
-    args = getntkeys(argstype)
-    data = getntkeys(datatype)
-    pars = getntkeys(parstype)
+    args = schema(argstype)
+    data = schema(datatype)
+    pars = schema(parstype)
 
     loader = @q begin
 
     end
 
-    for k in args
-        push!(loader.args, :($k = _args.$k))
+    for k in keys(args)
+        T = getproperty(args, k)
+        push!(loader.args, :($k::$T = _args.$k))
     end
-    for k in setdiff(data, pars)
-        push!(loader.args, :($k = _data.$k))
+    for k in setdiff(keys(data), keys(pars))
+        T = getproperty(data, k)
+        push!(loader.args, :($k::$T = _data.$k))
     end
 
-    for k in setdiff(pars, data)
-        push!(loader.args, :($k = _pars.$k))
+    for k in setdiff(keys(pars), keys(data))
+        T = getproperty(pars, k)
+        push!(loader.args, :($k::$T = _pars.$k))
     end
 
-    for k in pars ∩ data
-        push!(loader.args, :($k = Soss.NestedTuples.lazymerge(_data.$k, _pars.$k)))
+    for k in keys(pars) ∩ keys(data)
+        qk = QuoteNode(k)
+        pars_k = getproperty(pars, k)
+        data_k = getproperty(data, k)
+        if pars_k isa NamedTuple && data_k isa NamedTuple
+            push!(loader.args, :($k = Soss.NestedTuples.lazymerge(_data.$k, _pars.$k)))
+        else
+            T = getproperty(pars, k)
+            push!(loader.args, quote
+                _k = $qk
+                @warn "Duplicate key, ignoring $_k in data"
+                $k::$T = _pars.$k
+            end)
+        end
     end
 
     src -> (@q begin
@@ -235,7 +250,7 @@ end
 getntkeys(::NamedTuple{A,B}) where {A,B} = A
 getntkeys(::Type{NamedTuple{A,B}}) where {A,B} = A
 getntkeys(::Type{NamedTuple{A}}) where {A} = A
-getntkeys(::Type{LazyMerge{A,B,S,T}}) where {A,B,S,T} = Tuple(A ∪ B)
+getntkeys(::Type{LazyMerge{X,Y}}) where {X,Y} = union(getntkeys(X) ∪ getntkeys(Y))
 
 
 # These macros quickly define additional methods for when you get tired of typing `NamedTuple()`
diff --git a/src/distributions/dist.jl b/src/distributions/dist.jl
index 79b82e62..d86b7c50 100644
--- a/src/distributions/dist.jl
+++ b/src/distributions/dist.jl
@@ -1,7 +1,7 @@
-xform(d::EqualMix, _data) = xform(d.components[1], _data)
+as(d::EqualMix, _data) = as(d.components[1], _data)
 
 
 # StudentT(ν, μ = 0.0, σ = 1.0) = LocationScale(μ, σ, TDist(ν))
 
 
-xform(d::Dists.Dirichlet, _data) = UnitSimplex(length(d.alpha))
+as(d::Dists.Dirichlet, _data) = UnitSimplex(length(d.alpha))
diff --git a/src/distributions/iid.jl b/src/distributions/iid.jl
index e2ed6cd5..38dc9b28 100644
--- a/src/distributions/iid.jl
+++ b/src/distributions/iid.jl
@@ -1,8 +1,8 @@
 export iid
 
 
-iid(n::Int...) = dist -> iid(dist, n...)
+iid(n::Integer...) = dist -> iid(dist, n...)
 
-iid(dist::AbstractMeasure, n...) = dist ^ n
+iid(dist, n...) = dist ^ n
 
-iid(dist::Dists.Distribution, n) = Dists.Product(fill(dist, n))
+iid(dist::Dists.Distribution, n...) = MeasureBase.powermeasure(dist, n)
diff --git a/src/exponentialfamilies.jl b/src/exponentialfamilies.jl
index 5625e43f..7d0bda5c 100644
--- a/src/exponentialfamilies.jl
+++ b/src/exponentialfamilies.jl
@@ -20,7 +20,7 @@ end
 # This needs to be reworked, since many distributions are not from exponential families
 efam(d::ExpFamDist{P,X} where {P,X}) = d
 
-Distributions.logdensity(d::ExpFamDist{P,X}, x :: X) where {P,X} = d.fam.logh(x) + d.fam.logg(d.θ) + sum(d.fam.η(d.θ) .* d.fam.t(x))
+Distributions.logdensity_def(d::ExpFamDist{P,X}, x :: X) where {P,X} = d.fam.logh(x) + d.fam.logg(d.θ) + sum(d.fam.η(d.θ) .* d.fam.t(x))
 
 # Exponential families are closed under iid
 function iid(n::Integer, d::ExpFamDist{P,X} ) where {P,X}
@@ -64,7 +64,7 @@ binom = iid(1000,Bernoulli(0.2))
 
 using BenchmarkTools
 
-@btime logdensity(binom, repeat([true],1000))
+@btime logdensity_def(binom, repeat([true],1000))
 
 @btime sum(logdensity.(Bernoulli(0.2),repeat([true],1000)))
 
@@ -88,7 +88,7 @@ end
 
 const xs = 3 .+ 4 .* randn(1000);
 
-@btime logdensity(iid(1000,Normal(3.1,4.3)), xs)
+@btime logdensity_def(iid(1000,Normal(3.1,4.3)), xs)
 
 @btime sum(logdensity.(Normal(3,4),xs))
 
diff --git a/src/importance.jl b/src/importance.jl
index bce5b164..43fa3454 100644
--- a/src/importance.jl
+++ b/src/importance.jl
@@ -39,7 +39,7 @@ function sourceImportanceSample(_data)
         _datakeys = getntkeys(_data)
 
         function proc(m, st::Sample)
-            st.x ∈ _datakeys && return :(_ℓ += logdensity($(st.rhs), $(st.x)))
+            st.x ∈ _datakeys && return :(_ℓ += logdensity_def($(st.rhs), $(st.x)))
 
             if hasproperty(p.dists, st.x)
                 pdist = getproperty(p.dists, st.x)
@@ -81,7 +81,7 @@ end
 
 @inline function importanceSample(p, q, _data)
     x = rand(q)
-    ℓ = logdensity(p,x) - logdensity(q,x)
+    ℓ = logdensity_def(p,x) - logdensity_def(q,x)
     Weighted(ℓ,x)
 end
 
@@ -96,7 +96,7 @@ end
 #     # This determines how to initialize a Particle for a given expression
 #     vars(expr) = (bound(p) ∪ bound(q) ∪ stochastic(p) ∪ stochastic(q)) ∩ variables(expr)
 
-#     procp(p, st::Follows) = :($ℓ += logdensity($(st.rhs), $(st.x)))
+#     procp(p, st::Follows) = :($ℓ += logdensity_def($(st.rhs), $(st.x)))
 #     procp(p, st::Let)     = convert(Expr, st)
 #     procp(p, st::Return)  = nothing
 #     procp(p, st::LineNumber) = convert(Expr, st)
@@ -105,12 +105,12 @@ end
 #         if isempty(vars(st.rhs))
 #             @q begin
 #                 $(st.x) = Particles($N, $(st.rhs))
-#                 $ℓ -= logdensity($(st.rhs), $(st.x))
+#                 $ℓ -= logdensity_def($(st.rhs), $(st.x))
 #             end
 #         else
 #             @q begin
 #                 $(st.x) = rand($(st.rhs))
-#                 $ℓ -= logdensity($(st.rhs), $(st.x))
+#                 $ℓ -= logdensity_def($(st.rhs), $(st.x))
 #             end
 #         end
 #     end
diff --git a/src/inference/advancedhmc.jl b/src/inference/advancedhmc.jl
index 3a98fbbd..baf001a4 100644
--- a/src/inference/advancedhmc.jl
+++ b/src/inference/advancedhmc.jl
@@ -51,9 +51,9 @@ println("Posterior mean β: " * string(round(E_β, digits=2)))
 function advancedHMC(m :: ConditionalModel{A,B}, _data, N = 1000;
                                                          n_adapts  = 1000) where {A,B}
 
-    ℓ(pars) = logdensity(m, merge(pars, _data))
+    ℓ(pars) = logdensity_def(m, merge(pars, _data))
 
-    t = xform(m,_data)
+    t = as(m,_data)
 
     function f(x)
         (θ, logjac) = transform_and_logjac(t,x)
diff --git a/src/inference/dynamicHMC.jl b/src/inference/dynamicHMC.jl
index 2c7e3261..4b32ff7b 100644
--- a/src/inference/dynamicHMC.jl
+++ b/src/inference/dynamicHMC.jl
@@ -69,7 +69,7 @@ function dynamicHMC(
     rng::AbstractRNG,
     m::ConditionalModel,
     N::Int = 1000;
-    # method = logpdf,
+    # method = logdensityof,
     ad_backend = Val(:ForwardDiff),
     reporter = DynamicHMC.NoProgressReport(),
     kwargs...,
@@ -80,7 +80,7 @@ function dynamicHMC(
     ℓ = if haskey(kwargs, :ℓ)
         codegen(m; ℓ = kwargs[:ℓ])
     else 
-        (a, o, pars) -> _logdensity(M, Model(m), a, o, pars)
+        (a, o, pars) -> _logdensity_def(M, Model(m), a, o, pars)
     end
 
     _argvals = argvals(m)
@@ -88,7 +88,7 @@ function dynamicHMC(
 
     logp(pars) = ℓ(_argvals, _obs, pars)
 
-    t = xform(m)
+    t = as(m)
     P = LogDensityProblems.TransformedLogDensity(t, logp)
     ∇P = LogDensityProblems.ADgradient(ad_backend, P)
 
@@ -121,8 +121,8 @@ function dynamicHMC(
     kwargs...,
 )
     _data = m.obs
-    ℓ(pars) = logdensity(m, merge(pars, _data), method)
-    t = xform(m, _data)
+    ℓ(pars) = logdensity_def(m, merge(pars, _data), method)
+    t = as(m, _data)
     P = LogDensityProblems.TransformedLogDensity(t, ℓ)
     ∇P = LogDensityProblems.ADgradient(ad_backend, P)
 
@@ -155,7 +155,7 @@ end
 #     m::ConditionalModel,
 #     _data::NamedTuple,
 # )
-#     t = xform(m, _data)
+#     t = as(m, _data)
 #     (results, steps) = dynamicHMC(rng, m, _data, Val(Inf))
 #     Q = results.final_warmup_state.Q
 #     while true
diff --git a/src/inference/power-posterior.jl b/src/inference/power-posterior.jl
index 4dfb1560..b021e7ce 100644
--- a/src/inference/power-posterior.jl
+++ b/src/inference/power-posterior.jl
@@ -7,7 +7,7 @@ function power_posterior(cm::ConditionalModel)
     avs = argvals(cm)
     obs = observations(cm)
     function f(t, x) 
-        logdensity(pr(avs), x) + t * logdensity(lik(merge(avs, x)) | obs)
+        logdensity_def(pr(avs), x) + t * logdensity_def(lik(merge(avs, x)) | obs)
     end
 
     return f
diff --git a/src/inference/rejection.jl b/src/inference/rejection.jl
index 3a6a21ee..c92777ab 100644
--- a/src/inference/rejection.jl
+++ b/src/inference/rejection.jl
@@ -1,12 +1,12 @@
 
 
-function rejection(m)
-    r = makeRand(m)
-    function f(data)
-        prop = r(;data...)
-        while NamedTupleTools.select(prop, data) != data
-            prop = r(;data...)
-        end
-        return prop
-    end
-end
+# function rejection(m)
+#     r = makeRand(m)
+#     function f(data)
+#         prop = r(;data...)
+#         while NamedTupleTools.select(prop, data) != data
+#             prop = r(;data...)
+#         end
+#         return prop
+#     end
+# end
diff --git a/src/optim.jl b/src/optim.jl
index aa7138e3..709d5fc2 100644
--- a/src/optim.jl
+++ b/src/optim.jl
@@ -6,7 +6,7 @@ function makeLoss(model)
 
     t = getTransform(model)
 
-    fpre = @eval $(logdensity(model))
+    fpre = @eval $(logdensity_def(model))
     f(par, data) = Base.invokelatest(fpre, par, data)
 
     loss(x, data) = -f(t(x), data)
diff --git a/src/primitives/xform.jl b/src/primitives/as.jl
similarity index 64%
rename from src/primitives/xform.jl
rename to src/primitives/as.jl
index 143dc228..1b4bbb3a 100644
--- a/src/primitives/xform.jl
+++ b/src/primitives/as.jl
@@ -10,19 +10,19 @@ function NestedTuples.schema(::Type{TransformVariables.TransformTuple{T}}) where
 end
 
 # In Bijectors.jl,
-# logdensity_with_trans(dist, x, true) == logdensity(transformed(dist), link(dist, x))
+# logdensity_with_trans(dist, x, true) == logdensity_def(transformed(dist), link(dist, x))
 
 
-export xform
+export as
 
-xform(m::ConditionalModel{A, B}, _data::NamedTuple) where {A,B} = xform(m | _data)
+@inline as(m::ConditionalModel{A, B}, _data::NamedTuple) where {A,B} = as(m | _data)
 
-function xform(m::ConditionalModel{A, B}) where {A,B}
-    return _xform(getmoduletypencoding(m), Model(m), argvals(m), observations(m))
+@inline function as(m::ConditionalModel{A, B}) where {A,B}
+    return _as(getmoduletypencoding(m), Model(m), argvals(m), observations(m))
 end
 
-# function xform(m::Model{EmptyNTtype, B}) where {B}
-#     return xform(m,NamedTuple())    
+# function as(m::Model{EmptyNTtype, B}) where {B}
+#     return as(m,NamedTuple())    
 # end
 
 
@@ -44,7 +44,7 @@ function sourceXform(_data=NamedTuple())
             rhs = st.rhs
             
             thecode = @q begin 
-                _t = Soss.xform($rhs, get(_data, $xname, NamedTuple()))
+                _t = Soss.as($rhs, get(_data, $xname, NamedTuple()))
                 if !isnothing(_t)
                     _result = merge(_result, ($x=_t,))
                 end
@@ -71,15 +71,15 @@ end
 
 using Distributions: support
 
-function xform(d, _data::NamedTuple)
+@inline function as(d, _data::NamedTuple)
     if hasmethod(support, (typeof(d),))
         return asTransform(support(d)) 
     end
 
-    error("Not implemented:\nxform($d)")
+    error("Not implemented:\nas($d)")
 end
 
-using TransformVariables: ShiftedExp, ScaledShiftedLogistic, as
+using TransformVariables: ShiftedExp, ScaledShiftedLogistic
 
 function asTransform(supp:: Dists.RealInterval) 
     (lb, ub) = (supp.lb, supp.ub)
@@ -90,28 +90,28 @@ function asTransform(supp:: Dists.RealInterval)
     return ScaledShiftedLogistic(ub-lb, lb)
 end
 
-xform(d, _data) = nothing
+as(d, _data) = nothing
 
-xform(μ::AbstractMeasure,  _data::NamedTuple=NamedTuple()) = as(μ)
+as(μ::AbstractMeasure,  _data::NamedTuple) = as(μ)
 
-xform(d::Dists.AbstractMvNormal, _data::NamedTuple=NamedTuple()) = as(Array, size(d))
+as(d::Dists.AbstractMvNormal, _data::NamedTuple=NamedTuple()) = as(Array, size(d))
 
-@gg function _xform(M::Type{<:TypeLevel}, _m::Model{Asub,B}, _args::A, _data) where {Asub,A,B}
+@gg function _as(M::Type{<:TypeLevel}, _m::Model{Asub,B}, _args::A, _data) where {Asub,A,B}
     body = type2model(_m) |> sourceXform(_data) |> loadvals(_args, _data)
     @under_global from_type(_unwrap_type(M)) @q let M
         $body
     end    
 end
 
-function xform(d::Dists.Distribution{Dists.Univariate}, _data::NamedTuple=NamedTuple())
+function as(d::Dists.Distribution{Dists.Univariate}, _data::NamedTuple=NamedTuple())
     sup = Dists.support(d)
     lo = isinf(sup.lb) ? -TV.∞ : sup.lb
     hi = isinf(sup.ub) ? TV.∞ : sup.ub
     as(Real, lo,hi)
 end
 
-function xform(d::Dists.Product, _data::NamedTuple=NamedTuple())
+function as(d::Dists.Product, _data::NamedTuple=NamedTuple())
     n = length(d)
     v = d.v
-    as(Vector, xform(v[1]), n)
+    as(Vector, as(v[1]), n)
 end
diff --git a/src/primitives/basemeasure.jl b/src/primitives/basemeasure.jl
index 0942e5f7..0672fcf8 100644
--- a/src/primitives/basemeasure.jl
+++ b/src/primitives/basemeasure.jl
@@ -4,7 +4,7 @@ export basemeasure
 import MeasureTheory
 
 function MeasureTheory.basemeasure(c::ConditionalModel{A,B,M}, x=NamedTuple()) where {A,B,M}
-    _basemeasure(M, Model(c), argvals(c), observations(c), x)
+    _basemeasure(M, Model(c), argvals(c), x)  | observations(c)
 end
 
 export sourceBasemeasure
@@ -29,7 +29,7 @@ function sourceBasemeasure()
         wrap(kernel) = @q begin
             _bm = (;)
             $kernel
-            return Soss.productmeasure(identity, _bm)
+            return Soss.productmeasure(_bm)
         end
 
         buildSource(_m, proc, wrap) |> MacroTools.flatten
@@ -37,9 +37,9 @@ function sourceBasemeasure()
 end
 
 
-@gg function _basemeasure(M::Type{<:TypeLevel}, _m::Model, _args, _data, _pars)
-    body = type2model(_m) |> sourceBasemeasure() |> loadvals(_args, _data, _pars)
+@gg function _basemeasure(M::Type{<:TypeLevel}, _m::Model, _args, _pars)
+    body = type2model(_m) |> sourceBasemeasure() |> loadvals(_args, _pars)
     @under_global from_type(_unwrap_type(M)) @q let M
         $body
-    end
+    end 
 end
diff --git a/src/primitives/bijectors.jl b/src/primitives/bijectors.jl
index 9925ecb9..b4041efe 100644
--- a/src/primitives/bijectors.jl
+++ b/src/primitives/bijectors.jl
@@ -16,7 +16,7 @@ function (b::Inverse{<: Transform{N,T}})(y::AbstractVector) where {N,T <:  Trans
     b.orig.t(y)
 end
 
-Bijectors.bijector(d::ConditionalModel) = Transform(xform(d))
+Bijectors.bijector(d::ConditionalModel) = Transform(as(d))
 
 function Bijectors.logabsdetjac(b::Transform, x::NamedTuple)
     -transform_and_logjac(b.t, b.t(x))[2]
diff --git a/src/primitives/insupport.jl b/src/primitives/insupport.jl
new file mode 100644
index 00000000..07d168bc
--- /dev/null
+++ b/src/primitives/insupport.jl
@@ -0,0 +1,49 @@
+
+export logdensityof
+
+using NestedTuples: lazymerge
+import MeasureTheory
+
+import MeasureBase: insupport
+export insupport
+
+function MeasureBase.insupport(c::ConditionalModel{A,B,M}, x=NamedTuple()) where {A,B,M}
+    _insupport(M, Model(c), argvals(c), observations(c), x)
+end
+
+export sourceInsupport
+
+sourceInsupport(m::AbstractModel) = sourceInsupport()(Model(m))
+
+function sourceInsupport()
+    function(_m::Model)
+        proc(_m, st :: Assign)     = :($(st.x) = $(st.rhs))
+        proc(_m, st :: Return)     = nothing
+        proc(_m, st :: LineNumber) = nothing
+        function proc(_m, st :: Sample)
+            x = st.x
+            rhs = st.rhs
+            @q begin
+                Soss.dynamic(Soss.insupport($rhs, $x)) || return false
+                $x = Soss.predict($rhs, $x)
+            end
+        end
+
+        wrap(kernel) = @q begin
+            $kernel
+            return true
+        end
+
+        buildSource(_m, proc, wrap) |> MacroTools.flatten
+    end
+end
+
+# MeasureTheory.insupport(d::Distribution, val, tr) = logdensityof(d, val)
+
+
+@gg function _insupport(M::Type{<:TypeLevel}, _m::Model, _args, _data, _pars)
+    body = type2model(_m) |> sourceInsupport() |> loadvals(_args, _data, _pars)
+    @under_global from_type(_unwrap_type(M)) @q let M
+        $body
+    end
+end
diff --git a/src/primitives/likelihood-weighting.jl b/src/primitives/likelihood-weighting.jl
index 56e4d667..5f0b3f82 100644
--- a/src/primitives/likelihood-weighting.jl
+++ b/src/primitives/likelihood-weighting.jl
@@ -18,7 +18,7 @@ function sourceWeightedSample(_data)
         proc(_m, st :: LineNumber) = nothing
 
         function proc(_m, st :: Sample)
-            st.x ∈ _datakeys && return :(_ℓ += logdensity($(st.rhs), $(st.x)))
+            st.x ∈ _datakeys && return :(_ℓ += logdensity_def($(st.rhs), $(st.x)))
             return :($(st.x) = rand($(st.rhs)))
         end
 
diff --git a/src/primitives/logdensity.jl b/src/primitives/logdensity.jl
index eb4a03ed..81fa4cd2 100644
--- a/src/primitives/logdensity.jl
+++ b/src/primitives/logdensity.jl
@@ -1,28 +1,27 @@
 
-export logdensity
+export logdensityof
 
 using NestedTuples: lazymerge
 import MeasureTheory
 
-function MeasureTheory.logdensity(c::ConditionalModel{A,B,M}, x=NamedTuple()) where {A,B,M}
-    _logdensity(M, Model(c), argvals(c), observations(c), x)
+function MeasureBase.logdensityof(c::ConditionalModel{A,B,M}, x=NamedTuple()) where {A,B,M}
+    _logdensityof(M, Model(c), argvals(c), observations(c), x)
 end
 
-export sourceLogdensity
+export sourceLogdensityOf
 
-sourceLogdensity(m::AbstractModel) = sourceLogdensity()(Model(m))
+sourceLogdensityOf(m::AbstractModel) = sourceLogdensityOf()(Model(m))
 
-function sourceLogdensity()
+function sourceLogdensityOf()
     function(_m::Model)
         proc(_m, st :: Assign)     = :($(st.x) = $(st.rhs))
-        # proc(_m, st :: Sample)     = :(_ℓ += logdensity($(st.rhs), $(st.x)))
         proc(_m, st :: Return)     = nothing
         proc(_m, st :: LineNumber) = nothing
         function proc(_m, st :: Sample)
             x = st.x
             rhs = st.rhs
             @q begin
-                _ℓ += Soss.logdensity($rhs, $x)
+                _ℓ += Soss.logdensityof($rhs, $x)
                 $x = Soss.predict($rhs, $x)
             end
         end
@@ -37,56 +36,10 @@ function sourceLogdensity()
     end
 end
 
-# MeasureTheory.logdensity(d::Distribution, val, tr) = logpdf(d, val)
-
-
-@gg function _logdensity(M::Type{<:TypeLevel}, _m::Model, _args, _data, _pars)
-    body = type2model(_m) |> sourceLogdensity() |> loadvals(_args, _data, _pars)
+@gg function _logdensityof(M::Type{<:TypeLevel}, _m::Model, _args, _data, _pars)
+    body = type2model(_m) |> sourceLogdensityOf() |> loadvals(_args, _data, _pars)
     @under_global from_type(_unwrap_type(M)) @q let M
         $body
     end
 end
 
-
-function MeasureTheory.logpdf(c::ConditionalModel{A,B,M}, x=NamedTuple()) where {A,B,M}
-    _logpdf(M, Model(c), argvals(c), observations(c), x)
-end
-
-export sourceLogpdf
-
-sourceLogpdf(m::AbstractModel) = sourceLogpdf()(Model(m))
-
-function sourceLogpdf()
-    function(_m::Model)
-        proc(_m, st :: Assign)     = :($(st.x) = $(st.rhs))
-        # proc(_m, st :: Sample)     = :(_ℓ += logpdf($(st.rhs), $(st.x)))
-        proc(_m, st :: Return)     = nothing
-        proc(_m, st :: LineNumber) = nothing
-        function proc(_m, st :: Sample)
-            x = st.x
-            rhs = st.rhs
-            @q begin
-                _ℓ += Soss.logpdf($rhs, $x)
-                $x = Soss.predict($rhs, $x)
-            end
-        end
-
-        wrap(kernel) = @q begin
-            _ℓ = 0.0
-            $kernel
-            return _ℓ
-        end
-
-        buildSource(_m, proc, wrap) |> MacroTools.flatten
-    end
-end
-
-# MeasureTheory.logpdf(d::Distribution, val, tr) = logpdf(d, val)
-
-
-@gg function _logpdf(M::Type{<:TypeLevel}, _m::Model, _args, _data, _pars)
-    body = type2model(_m) |> sourceLogpdf() |> loadvals(_args, _data, _pars)
-    @under_global from_type(_unwrap_type(M)) @q let M
-        $body
-    end
-end
diff --git a/src/samplechains/dynamichmc.jl b/src/samplechains/dynamichmc.jl
index c3985860..7a126f5c 100644
--- a/src/samplechains/dynamichmc.jl
+++ b/src/samplechains/dynamichmc.jl
@@ -13,8 +13,8 @@ function sample(rng::AbstractRNG,
     nsamples::Int=1000,
     nchains::Int=4)
 
-    ℓ(x) = MeasureTheory.logpdf(m, x)
-    tr = xform(m)
+    ℓ(x) = MeasureTheory.logdensityof(m, x)
+    tr = as(m)
 
     chains = newchain(rng, nchains, config, ℓ, tr)
     sample!(chains, nsamples - 1)
diff --git a/src/scratch/bouncy.jl b/src/scratch/bouncy.jl
index 12330479..c90563ae 100644
--- a/src/scratch/bouncy.jl
+++ b/src/scratch/bouncy.jl
@@ -17,9 +17,9 @@ Returns a `Trace` object.
 """
 function bouncy(m::ConditionalModel, T = 1000.0;  c=10.0, λref=0.1, ρ=0.0, adapt=false) where {A,B}
 
-    ℓ(pars) = logdensity(m, pars)
+    ℓ(pars) = logdensity_def(m, pars)
 
-    t = xform(m)
+    t = as(m)
 
     function f(x)
         (θ, logjac) = transform_and_logjac(t, x)
diff --git a/src/scratch/zigzag.jl b/src/scratch/zigzag.jl
index f4f652bc..4659542d 100644
--- a/src/scratch/zigzag.jl
+++ b/src/scratch/zigzag.jl
@@ -5,7 +5,7 @@ function zigzag(m::ConditionalModel, T = 1000.0; c=10.0, adapt=false) where {A,B
 
     ℓ = Base.Fix1(logdensity, m)
 
-    t = xform(m)
+    t = as(m)
 
     function f(x)
         (θ, logjac) = transform_and_logjac(t, x)
diff --git a/src/symbolic/symbolic.jl b/src/symbolic/symbolic.jl
index 06a2e083..d3526c08 100644
--- a/src/symbolic/symbolic.jl
+++ b/src/symbolic/symbolic.jl
@@ -19,9 +19,9 @@ export schema
 
 export symlogdensity
 
-symlogdensity(d, x::Symbolic) = logpdf(d,x)
+symlogdensity(d, x::Symbolic) = logdensityof(d,x)
 
-function symlogdensity(d::ProductMeasure{F,S,<:AbstractArray}, x::Symbolic{A}) where {F,S,A <: AbstractArray}
+function symlogdensity(d::ProductMeasure, x::Symbolic{A}) where {A <: AbstractArray}
     dims = size(d)
 
     iters = Sym{Int}.(gensym.(Symbol.(:i, 1:length(dims))))
diff --git a/src/transforms/utils.jl b/src/transforms/utils.jl
index d9872401..180ca10d 100644
--- a/src/transforms/utils.jl
+++ b/src/transforms/utils.jl
@@ -125,7 +125,8 @@ function after(m::Model, xs...; strict = false)
     parms = setdiff(parms, args)
     trim_args!(args, m, parms)
     m1 = setReturn(assemblefrom(m, parms, args), getReturn(m))
-    newargs = arguments(m1) ∪ setdiff(variables(m) ∩ variables(getReturn(m1)), variables(m1))
+    newargs = Symbol[]
+    append!(newargs, arguments(m1) ∪ setdiff(variables(m) ∩ variables(getReturn(m1)), variables(m1)))
     return merge(m1, Model(getmodule(m), newargs, quote end))
 end
 
diff --git a/test/linear-model.jl b/test/linear-model.jl
index f8292c19..d834c4a9 100644
--- a/test/linear-model.jl
+++ b/test/linear-model.jl
@@ -44,6 +44,6 @@ jointdist = m2(X=X)
 
 @testset "logdensity, symlogdensity" begin
     # @test_nowarn symlogdensity(m2).evalf(3)
-    @test logdensity(jointdist, (β = truth.trace.β, y=y)) ≈ -8.015157812948065
-    # @test logdensity(jointdist, truth, codegen) ≈ -28.551921801470904
+    @test logdensity_def(jointdist, (β = truth.trace.β, y=y)) ≈ -8.015157812948065
+    # @test logdensity_def(jointdist, truth, codegen) ≈ -28.551921801470904
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 98bd47df..b12ba398 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,8 @@
 using Soss
 using Test
 using MeasureTheory
-using TransformVariables
+import TransformVariables as TV
+using TransformVariables: transform
 using Aqua
 Aqua.test_all(Soss; ambiguities=false, unbound_args=false)
 
@@ -48,8 +49,8 @@ include("examples-list.jl")
 
         x = rand(outer(sub=inner)).m
         post = outer(sub=inner) | (m = (x=x,),)
-        t = xform(post)
-        @test logdensity(post, transform(t, randn(3))) isa Float64
+        t = as(post)
+        @test logdensityof(post, transform(t, randn(3))) isa Real
     end
 
     @testset "Predict" begin
@@ -59,7 +60,7 @@ include("examples-list.jl")
             return y
         end
             
-        mean(predict(m(), [(p=p,) for p in rand(10000)])) isa Float64
+        @test mean(predict(m(), [(p=p,) for p in rand(10000)])) isa Float64
     end
 
     @testset "https://github.com/cscherrer/Soss.jl/issues/258" begin
@@ -76,7 +77,7 @@ include("examples-list.jl")
 
         mm = m2(m=m1())
         
-        @test xform(mm|(y=1.0,)) isa TransformVariables.TransformTuple
+        @test as(mm|(y=1.0,)) isa TV.TransformTuple
         @test basemeasure(mm | (y=1.0,)) isa ProductMeasure
         @test testvalue(mm) isa NamedTuple
     end
@@ -97,7 +98,7 @@ include("examples-list.jl")
         
         mm = m2(m=m1())
 
-        @test xform(mm|(;y=1.0,)) isa TransformVariables.TransformTuple
+        @test as(mm|(;y=1.0,)) isa TV.TransformTuple
         @test basemeasure(mm | (y=1.0,)) isa ProductMeasure
         @test testvalue(mm) isa NamedTuple
     end
@@ -107,7 +108,7 @@ include("examples-list.jl")
             x ~ For(3) do j Normal(μ=j) end
         end;
         
-        @test logpdf(m(), rand(m())) isa Float64
+        @test logdensityof(m(), rand(m())) isa Float64
     end
 
     @testset "Local variables" begin
@@ -138,7 +139,7 @@ include("examples-list.jl")
 
         post = m() | (c=c,)
 
-        @test transform(xform(post), randn(6)) isa NamedTuple
+        @test transform(as(post), randn(6)) isa NamedTuple
 
         @testset "logdensity" begin
             dat = randn(100)
@@ -151,11 +152,20 @@ include("examples-list.jl")
             mod = m( (; n = length(dat) ) )
             post = mod | (data = dat,)
 
-            @test logdensity( mod( (μ = 1., σ = 2., data = dat) ) ) == logdensity( post( (μ = 1., σ = 2.) ) )
+            @test logdensityof( mod( (μ = 1., σ = 2., data = dat) ) ) == logdensityof( post( (μ = 1., σ = 2.) ) )
         end
 
 
     end
 
+    @testset "basemeasure" begin
+        m = @model n begin
+            p ~ Uniform()
+            x ~ Bernoulli(p) ^ n
+        end
 
+        post = m(10) | (x = rand(Bool, 10),)
+        base = basemeasure(post)
+        @test logdensity_def(base, (p=0.2, x=post.obs.x)) isa Real
+    end
 end