Fixed transformations

HISTORY.md

@@ -1,6 +1,90 @@
 # 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.
 

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.

docs/src/api.md

@@ -440,19 +440,12 @@ DynamicPPL.setindex_internal!!
 
 #### Transformations
 
-```@docs
-DynamicPPL.AbstractTransformation
-DynamicPPL.NoTransformation
-DynamicPPL.DynamicTransformation
-DynamicPPL.StaticTransformation
-```
-
 ```@docs
 DynamicPPL.link
 DynamicPPL.invlink
 DynamicPPL.link!!
 DynamicPPL.invlink!!
-DynamicPPL.update_link_status!!
+DynamicPPL.update_transform_status!!
 ```
 
 ```@docs
@@ -461,21 +454,19 @@ DynamicPPL.LinkAll
 DynamicPPL.UnlinkAll
 DynamicPPL.LinkSome
 DynamicPPL.UnlinkSome
+DynamicPPL.WithTransforms
 ```
 
 ```@docs
 DynamicPPL.AbstractTransform
 DynamicPPL.DynamicLink
 DynamicPPL.Unlink
+DynamicPPL.FixedTransform
+DynamicPPL.NoTransform
 DynamicPPL.target_transform
 DynamicPPL.apply_transform_strategy
 ```
 
-```@docs
-DynamicPPL.transformation
-DynamicPPL.default_transformation
-```
-
 #### Utils
 
 ```@docs
@@ -561,14 +552,10 @@ init
 get_param_eltype
 ```
 
-The function [`DynamicPPL.init`](@ref) should return an `AbstractTransformedValue`.
-There are three subtypes currently available:
+The function [`DynamicPPL.init`](@ref) should return a `TransformedValue`.
 
 ```@docs
-DynamicPPL.AbstractTransformedValue
-DynamicPPL.VectorValue
-DynamicPPL.LinkedVectorValue
-DynamicPPL.UntransformedValue
+DynamicPPL.TransformedValue
 ```
 
 The interface for working with transformed values consists of:

docs/src/evaluation.md

@@ -249,7 +249,7 @@ getlogprior(accs) - getlogjac(accs)
 
 You might ask: given that we specified parameters in untransformed space, how do we then retrieve the parameters in transformed space?
 The answer to this is to use an accumulator (no surprises there!) that collects the transformed values.
-Specifically, a `VectorValueAccumulator` collects vectorised forms of the parameters, which may either be [`VectorValue`](@ref)s or [`LinkedVectorValue`](@ref)s.
+Specifically, a `VectorValueAccumulator` collects vectorised forms of the parameters: that is, `TransformedValue{V,T}` where `V<:AbstractVector`.
 
 ```@example 1
 accs = OnlyAccsVarInfo(VectorValueAccumulator())

docs/src/init.md

@@ -29,12 +29,12 @@ The subsequent sections will demonstrate how this can be done.
 
 ## The required interface
 
-Each initialisation strategy must subtype `AbstractInitStrategy`, and implement `DynamicPPL.init(rng, vn, dist, strategy)`, which must return an `AbstractTransformedValue`.
+Each initialisation strategy must subtype `AbstractInitStrategy`, and implement `DynamicPPL.init(rng, vn, dist, strategy)`, which must return a `TransformedValue`.
 
 ```@docs; canonical=false
 AbstractInitStrategy
 init
-DynamicPPL.AbstractTransformedValue
+DynamicPPL.TransformedValue
 ```
 
 ## An example
@@ -64,7 +64,7 @@ function DynamicPPL.init(rng, vn::VarName, ::Distribution, strategy::InitRandomW
     new_x = rand(rng, Normal(strategy.x_prev, strategy.step_size))
     # Insert some printing to see when this is called.
     @info "init() is returning: $new_x"
-    return DynamicPPL.UntransformedValue(new_x)
+    return DynamicPPL.TransformedValue(new_x, DynamicPPL.NoTransform())
 end
 ```
 
@@ -103,19 +103,19 @@ In this case, we have defined an initialisation strategy that is random (and thu
 However, initialisation strategies can also be fully deterministic, in which case the `rng` argument is not needed.
 For example, [`DynamicPPL.InitFromParams`](@ref) reads from a set of parameters which are known ahead of time.
 
-## The returned `AbstractTransformedValue`
+## The returned `TransformedValue`
 
-As mentioned above, the `init` function must return an `AbstractTransformedValue`.
-The subtype of `AbstractTransformedValue` used does not affect the result of the model evaluation, but it may have performance implications.
+As mentioned above, the `init` function must return an `TransformedValue`.
+The transform stored inside this does not affect the result of the model evaluation, but it may have performance implications.
 **In particular, the returned subtype does not determine whether the log-Jacobian term is accumulated or not: that is determined by a separate [_transform strategy_](@ref transform-strategies).**
 
-What this means is that initialisation strategies should always choose the laziest possible subtype of `AbstractTransformedValue`.
+What this means is that initialisation strategies should always choose the laziest possible version of `TransformedValue`, electing to do as few transformations as possible inside `init`.
 
-For example, in the above example, we used `UntransformedValue`, which is the simplest possible choice.
+For example, in the above example, we simply wrapped the untransformed value in `TransformedValue(..., NoTransform())`, which is the simplest possible choice.
 If a linked value is required by a later step inside `tilde_assume!!` (either the transformation or accumulation steps), it is the responsibility of that step to perform the linking.
 
 Conversely, [`DynamicPPL.InitFromUniform`](@ref) samples inside linked space.
-Instead of performing the inverse link transform and returning an `UntransformedValue`, it directly returns a `LinkedVectorValue`: this means that if a linked value is required by a later step, it is not necessary to link it again.
+Instead of performing the inverse link transform eagerly, it directly returns a `TransformedValue(val, DynamicLink())`, where `val` is *already* the linked vector: this means that if a linked value is required by a later step, it is not necessary to link it again.
 Even if no linked value is required, this lazy approach does not hurt performance, as it just defers the inverse linking to the later step.
 
 In both cases, only one linking operation is performed (at most).

docs/src/migration.md

@@ -36,7 +36,11 @@ Old:
 
 ```@example 1
 vi = VarInfo(Xoshiro(468), model)
-vi[@varname(x)], vi[@varname(y)]
+# This no longer works, but you may have used it.
+# vi[@varname(x)], vi[@varname(y)]
+
+# This still works
+DynamicPPL.getindex_internal(vi, @varname(x))
 ```
 
 New: