Fixed transformations

.github/workflows/Benchmarking.yml

@@ -16,7 +16,7 @@ jobs:
       - uses: julia-actions/setup-julia@v2
         with:
           version: '1.11'
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - name: Run benchmarks
         id: benchmark
@@ -41,7 +41,7 @@ jobs:
       - uses: julia-actions/setup-julia@v2
         with:
           version: '1.11'
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - name: Run benchmarks
         id: benchmark
@@ -64,7 +64,7 @@ jobs:
       - uses: julia-actions/setup-julia@v2
         with:
           version: '1.11'
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - name: Combine benchmark results
         working-directory: ./benchmarks

.github/workflows/CI.yml

@@ -62,7 +62,7 @@ jobs:
         with:
           version: ${{ matrix.runner.version }}
 
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - uses: julia-actions/julia-buildpkg@v1
 
@@ -75,7 +75,7 @@ jobs:
 
       - uses: julia-actions/julia-processcoverage@v1
 
-      - uses: codecov/codecov-action@v5
+      - uses: codecov/codecov-action@v6
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/DocTest.yml

@@ -32,7 +32,7 @@ jobs:
         with:
           version: '1'
 
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - uses: julia-actions/julia-buildpkg@v1
 

.github/workflows/Enzyme.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           version: "1.11"
 
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - name: Run AD with Enzyme on demo models
         working-directory: test/integration/enzyme

.github/workflows/FloatTypes.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           version: "1"
 
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
 
       - name: Run float type tests
         working-directory: test/floattypes

.github/workflows/JuliaPre.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: julia-actions/setup-julia@v2
         with:
           version: 'pre' # pre-release
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:

HISTORY.md

@@ -1,8 +1,105 @@
 # 0.41
 
-Removed the `varinfo` keyword argument from `DynamicPPL.TestUtils.AD.run_ad` and replaced the `varinfo` field in `ADResult` with `ldf::LogDensityFunction`.
+## Breaking changes
+
+### Unification of transformed values
+
+Previously, there were separate types `UntransformedValue`, `VectorValue`, and `LinkedVectorValue`, which were all subtypes of `AbstractTransformedValue`.
+The abstract type has been removed, and all of these have been unified in a single `TransformedValue` struct, which wraps the (maybe transformed) value, plus an `AbstractTransform` that describes the inverse transformation (to get back to the raw value).
+
+Concretely,
+
+  - `UntransformedValue(val)` is now `TransformedValue(val, NoTransform())`
+  - `VectorValue(vec, tfm)` is now `TransformedValue(vec, Unlink())`
+  - `LinkedVectorValue(vec, tfm)` is now `TransformedValue(vec, DynamicLink())`
+
+**Note that this means for `VectorValue` and `LinkedVectorValue`, the transform is no longer stored on the value itself.**
+This means that given one of these values, you *cannot* access the raw value without running the model.
+
+The reason why this is done is that the transform may in principle change between model executions.
+This can happen if the prior distribution of a variable depends on the value of another variable.
+Previously, in DynamicPPL, we *always* made sure to recompute the transform during model evaluation; however, this was not enforced by the data structure.
+The current implementation makes it impossible to accidentally use an outdated transform, and is therefore more robust.
+
+### Addition of `FixedTransform`
+
+The above unification allows us to introduce a new transform subtype, `FixedTransform{F}`, which wraps a known function `F` that is assumed to always be static, allowing the transform to be cached and reused across model executions.
+**This should only be used when it is known ahead of time that the transform will never change between model executions.**
+It is the user's responsibility to ensure that this is the case.
+Using `FixedTransform` when the transform does change between model executions can lead to incorrect results.
+
+For many simple distributions, this in fact saves absolutely no time, because deriving the transform from the distribution takes almost negligible time (~ 1 ns!).
+However, there are some edge cases for which this is not the case: for example, `product_distribution([Beta(2, 2), Normal()])` is quite slow (~ 3 µs).
+In such cases, using `FixedTransform` can lead to substantial performance improvements.
+
+To use `FixedTransform` with `LogDensityFunction`, you need to:
+
+ 1. Create a `VarNamedTuple` mapping `VarName`s to `FixedTransform`s for the variables in your model.
+    This can be done using `DynamicPPL.FixedTransformAccumulator` (see the DynamicPPL docs for more info), but is most easily done by calling `get_fixed_transforms(model, transform_strategy)`, where `transform_strategy` says whether you want linked or unlinked transforms.
+
+ 2. Wrap the `VarNamedTuple` inside `WithTransforms(vnt, UnlinkAll())`.
+    `WithTransforms` is a subtype of `AbstractTransformStrategy`, much like `LinkAll()`.
+    However, `WithTransforms` specifies that *these exact transforms are to be used*, whereas `LinkAll` says 'derive the transforms again at model runtime'.
+ 3. Construct a `LogDensityFunction(model, getlogjoint_internal, WithTransforms(...)); adtype=adtype`.
+
+### Removal of `getindex(vi::VarInfo, vn::VarName)`
+
+The main role of `VarInfo` was to store vectorised transformed values.
+Previously, these were stored as `VectorValue`s or `LinkedVectorValue`s: these used to carry the inverse transform with them, which allowed you to access the raw value via `vi[vn]`, or equivalently `getindex(vi, vn)`.
+
+The problem with this is that (as described above) the correct transform may depend on the values of the variables themselves.
+That means that if we update the vectorised value without changing the transform, we could end up with an inconsistent state and incorrect results.
+In particular, this is *exactly* what the function `unflatten!!` does: it updates the vectorised values but does not touch the transform.
+
+In the current version, we have removed this method to prevent the possibility of obtaining incorrect results.
+(Our hands are also forced by the fact that the new `TransformedValue`s do not store the actual transform with them.)
+
+*In place of using `VarInfo`, we strongly recommend that you migrate to using `OnlyAccsVarInfo`.*
+In particular, to access raw (untransformed) values, you should use an `OnlyAccsVarInfo` with a `RawValueAccumulator`.
+There is [a migration guide available on the DynamicPPL documentation](https://turinglang.org/DynamicPPL.jl/stable/migration/) and we are very happy to add more examples to this if you run into something that is not covered.
+
+### FixedTransformAccumulator
+
+TODO, this part is still being worked on.
+
+  - `BijectorAccumulator` → `FixedTransformAccumulator`
+  - `get_fixed_transforms(::VarInfo)`
+  - `get_fixed_transforms(::Model)`
+
+## Miscellaneous breaking changes
+
+  - Removed the `varinfo` keyword argument from `DynamicPPL.TestUtils.AD.run_ad`, and replaced the `varinfo` field in the returned `ADResult` with `ldf::LogDensityFunction`.
+
+## Internal changes
+
+The following functions were not exported; we document changes in them for completeness.
+
+  - Given the above changes, the old framework of `AbstractTransformation`, `StaticTransformation`, and `DynamicTransformation` are no longer needed, and have been removed.
+    The (rarely used) methods of `link`, `link!!`, `invlink`, and `invlink!!` that took an `AbstractTransformation` as the first argument have been removed.
+    The same is true of the functions `default_transformation` and `transformation`.
+
+  - `RangeAndLinked` has been expanded to `RangeAndTransform`: instead of just carrying a Boolean indicating whether the transform is `DynamicLink` or `Unlink`, it now stores the full transform.
+    This is done in order to accommodate `FixedTransform`.
+  - Consequently, `get_ranges_and_linked` has been renamed to `get_rangeandtransforms`.
+    Its function is still to return a `VarNamedTuple` of `RangeAndTransform`s.
+  - `update_link_status!!` has been renamed to `update_transform_status!!`.
+  - `get_transform_strategy` has been renamed to `infer_transform_strategy`.
+  - `from_internal_transform` and `from_linked_internal_transform` have been removed, since the new `TransformedValue`s do not store the transform with them.
+  - Removed `getargnames`, `getmissings`, and `Base.nameof(::Model)` from the public API (export and documentation) as they are considered internal implementation details.
+
+# 0.40.20
+
+Added a public function, `DynamicPPL.extract_prefixes(::AbstractContext)`, to more generally handle the removal of `PrefixContext` entries from the context stack.
+
+Also exports `SkipTemplate` and `NoTemplate` from DynamicPPL.
+
+# 0.40.19
+
+Fixed `check_model()` occasionally failing to catch duplicate `VarName` assignments correctly when using multithreading with thread-safe variables.
+
+# 0.40.18
 
-Removed `getargnames`, `getmissings`, and `Base.nameof(::Model)` from the public API (export and documentation) as they are considered internal implementation details.
+Added a check on `unflatten!!` to error if the input vector was too long.
 
 # 0.40.17
 

docs/Project.toml

@@ -1,9 +1,11 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

docs/make.jl

@@ -52,6 +52,7 @@ makedocs(;
         "Tilde-statements" => "tilde.md",
         "Initialisation strategies" => "init.md",
         "Transform strategies" => "transforms.md",
+        "Fixed transforms" => "transforms_fixed.md",
         "Accumulators" => [
             "accs/overview.md",
             "accs/existing.md",

docs/src/accs/existing.md

@@ -42,4 +42,6 @@ get_vector_params
 ```@docs
 PriorDistributionAccumulator
 get_priors
+FixedTransformAccumulator
+get_fixed_transforms
 ```

docs/src/accs/values.md

@@ -29,8 +29,10 @@ struct VarInfo{Tfm<:AbstractTransformStrategy,V<:VarNamedTuple,A<:AccumulatorTup
 end
 ```
 
-The `values` field stores either [`LinkedVectorValue`](@ref)s or [`VectorValue`](@ref)s.
-The `transform_strategy` field stores an `AbstractTransformStrategy` which is (as far as possible) consistent with the type of values stored in `values`.
+The `values` field stores `DynamicPPL.TransformedValue`s, but it is mandatory that these transformed values are vectorised.
+That is, it is permissible to store (for example) `TransformedValue([1.0], Unlink())`, but not `TransformedValue(1.0, NoTransform())`.
+
+Furthermore, the `transform_strategy` field stores an `AbstractTransformStrategy` which is (as far as possible) consistent with the type of values stored in `values`.
 
 Here is an example:
 
@@ -46,7 +48,7 @@ vi = VarInfo(dirichlet_model)
 vi
 ```
 
-In `VarInfo`, it is mandatory to store `LinkedVectorValue`s or `VectorValue`s as `ArrayLikeBlock`s (see the [Array-like blocks](@ref array-like-blocks) documentation for information on this).
+In `vi.values`, it is mandatory to store `TransformedValue`s as `ArrayLikeBlock`s (see the [Array-like blocks](@ref array-like-blocks) documentation for information on this).
 The reason is because, if the value is linked, it may have a different size than the number of indices in the `VarName`.
 This means that when retrieving the keys, we obtain each block as a single key:
 
@@ -60,21 +62,21 @@ In a `VarInfo`, the `accs` field is responsible for the accumulation step, just
 
 However, `values` serves three purposes in one:
 
-  - it is sometimes used for initialisation (when the model's leaf context is `DefaultContext`, the `AbstractTransformedValue` to be used in the transformation step is read from it)
+  - it is sometimes used for initialisation (when the model's leaf context is `DefaultContext`, the `TransformedValue` to be used in the transformation step is read from it)
   - it also determines whether the log-Jacobian term should be included or not (if the value is a `LinkedVectorValue`, the log-Jacobian is included)
-  - it is sometimes also used for accumulation (when evaluating a model with a VarInfo, we will potentially store a new `AbstractTransformedValue` in it!).
+  - it is sometimes also used for accumulation (when evaluating a model with a VarInfo, we will potentially store a new `TransformedValue` in it!).
 
 The path to removing `VarInfo` is essentially to separate these three roles:
 
  1. The initialisation role of `varinfo.values` can be taken over by an initialisation strategy that wraps it.
-    Recall that the only role of an initialisation strategy is to provide an `AbstractTransformedValue` via [`DynamicPPL.init`](@ref).
+    Recall that the only role of an initialisation strategy is to provide an `TransformedValue` via [`DynamicPPL.init`](@ref).
     This can be trivially done by indexing into the `VarNamedTuple` stored in the strategy.
 
  2. Whether the log-Jacobian term should be included or not can be determined by a transform strategy.
     Much like how we can have an initialisation strategy that takes values from a `VarInfo`, we can also have a transform strategy that is defined by the existing status of a `VarInfo`.
     This is implemented in the `DynamicPPL.get_link_strategy(::AbstractVarInfo)` function.
  3. The accumulation role of `varinfo.values` can be taken over by a new accumulator, which we call `VectorValueAccumulator`.
-    This name is chosen because it does not store generic `AbstractTransformedValue`s, but only two subtypes of it, `LinkedVectorValue` and `VectorValue`.
+    This name is chosen because it does not store generic `TransformedValue`s, but only two subtypes of it, `LinkedVectorValue` and `VectorValue`.
     `VectorValueAccumulator` is implemented inside `src/accs/vector_value.jl`.
 
 !!! note
@@ -86,11 +88,11 @@ The path to removing `VarInfo` is essentially to separate these three roles:
 
 ## `RawValueAccumulator`
 
-Earlier we said that `VectorValueAccumulator` stores only two subtypes of `AbstractTransformedValue`: `LinkedVectorValue` and `VectorValue`.
-One might therefore ask about the third subtype, namely, `UntransformedValue`.
+Earlier we said that `VectorValueAccumulator` stores only values that have been vectorised.
+One might therefore ask about unvectorised values — and in particular, values that have *not* been transformed at all, i.e., `TransformedValue(val, NoTransform())`.
 
-It turns out that it is very often useful to store [`UntransformedValue`](@ref)s.
-Additionally, since `UntransformedValue`s must always correspond exactly to the indices they are assigned to, we can unwrap them and do not need to store them as array-like blocks!
+It turns out that it is very often useful to store such untransformed values.
+Additionally, since the values must always correspond exactly to the indices they are assigned to, we can unwrap them and do not need to store them as array-like blocks!
 
 This is the role of `RawValueAccumulator`.
 
@@ -100,7 +102,7 @@ _, oavi = DynamicPPL.init!!(dirichlet_model, oavi, InitFromPrior(), UnlinkAll())
 raw_vals = get_raw_values(oavi)
 ```
 
-Note that when we unwrap `UntransformedValue`s, we also lose the block structure that was present in the model.
+Note that when we unwrap `TransformedValue`s, we also lose the block structure that was present in the model.
 That means that in `RawValueAccumulator`, there is no longer any notion that `x[1:3]` was set together, so the keys correspond to the individual indices.
 
 ```@example 1
@@ -116,24 +118,24 @@ This is why indices of keys like `x[1:3] ~ dist` end up being split up in chains
 
 ## Why do we still need to store `TransformedValue`s?
 
-Given that `RawValueAccumulator` exists, one may wonder why we still need to store the other `AbstractTransformedValue`s at all, i.e. what the purpose of `VectorValueAccumulator` is.
+Given that `RawValueAccumulator` exists, one may wonder why we still need to store the other `TransformedValue`s at all, i.e. what the purpose of `VectorValueAccumulator` is.
 
 Currently, the only remaining reason for transformed values is the fact that we may sometimes need to perform [`DynamicPPL.unflatten!!`](@ref) on a `VarInfo`, to insert new values into it from a vector.
 
 ```@example 1
 vi = VarInfo(dirichlet_model)
-vi[@varname(x[1:3])]
+DynamicPPL.getindex_internal(vi, @varname(x[1:3]))
 ```
 
 ```@example 1
 vi = DynamicPPL.unflatten!!(vi, [0.2, 0.5, 0.3])
-vi[@varname(x[1:3])]
+DynamicPPL.getindex_internal(vi, @varname(x[1:3]))
 ```
 
 If we do not store the vectorised form of the values, we will not know how many values to read from the input vector for each key.
 
 Removing upstream usage of `unflatten!!` would allow us to completely get rid of `TransformedValueAccumulator` and only ever use `RawValueAccumulator`.
 See [this DynamicPPL issue](https://github.com/TuringLang/DynamicPPL.jl/issues/836) for more information.
 
-One possibility for removing `unflatten!!` is to turn it into a function that, instead of generating a new VarInfo, instead generates a tuple of new initialisation and link strategies which returns `LinkedVectorValue`s or `VectorValue`s containing views into the input vector.
+One possibility for removing `unflatten!!` is to turn it into a function that, instead of generating a new VarInfo, instead generates a tuple of new initialisation and transform strategies similar to `InitFromVector`.
 This would be conceptually very similar to how `LogDensityFunction` currently works.