Julia Version 1.11.7
Commit f2b3dbda30a (2025-09-08 12:10 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)

┌───────────────────────┬───────┬─────────────┬───────────────────┬────────┬────────────────┬─────────────────┐
│                 Model │   Dim │  AD Backend │           VarInfo │ Linked │ t(eval)/t(ref) │ t(grad)/t(eval) │
├───────────────────────┼───────┼─────────────┼───────────────────┼────────┼────────────────┼─────────────────┤
│ Simple assume observe │     1 │ forwarddiff │             typed │  false │            6.5 │             1.8 │
│           Smorgasbord │   201 │ forwarddiff │             typed │  false │          732.4 │            44.4 │
│           Smorgasbord │   201 │ forwarddiff │ simple_namedtuple │   true │          411.5 │            58.9 │
│           Smorgasbord │   201 │ forwarddiff │           untyped │   true │          784.9 │            36.1 │
│           Smorgasbord │   201 │ forwarddiff │       simple_dict │   true │         7127.7 │            27.5 │
│           Smorgasbord │   201 │ reversediff │             typed │   true │          736.1 │            55.9 │
│           Smorgasbord │   201 │    mooncake │             typed │   true │          715.6 │             6.4 │
│           Smorgasbord │   201 │      enzyme │             typed │   true │          889.4 │             3.9 │
│    Loop univariate 1k │  1000 │    mooncake │             typed │   true │         3918.4 │             5.7 │
│       Multivariate 1k │  1000 │    mooncake │             typed │   true │         1007.5 │             8.9 │
│   Loop univariate 10k │ 10000 │    mooncake │             typed │   true │        42724.2 │             5.3 │
│      Multivariate 10k │ 10000 │    mooncake │             typed │   true │         8594.6 │            10.1 │
│               Dynamic │    10 │    mooncake │             typed │   true │          125.0 │            11.8 │
│              Submodel │     1 │    mooncake │             typed │   true │            9.0 │             6.4 │
│                   LDA │    12 │ reversediff │             typed │   true │          998.5 │             2.0 │
└───────────────────────┴───────┴─────────────┴───────────────────┴────────┴────────────────┴─────────────────┘

"""
DynamicPPL transforms
---------------------

In general, there are four ways that parameters are internally represented:

1. **Unlinked**: The default, unchanged, representation of the parameter, as returned by
   `rand(dist)`.

2. **Unlinked vectorised**: A flattened vector form of the original representation.

3. **Linked**: A transformed version of the original representation. This is accomplished
   by applying a bijective transformation (a 'link function') to the unlinked parameter.
   This results in a parameter that exists in unconstrained Euclidean space.

4. **Linked vectorised**: A flattened vector form of the linked representation.

## Example 1

Consider the variable `x ~ LogNormal()`.

1. If we draw a sample from this distribution, it is necessarily positive, so the unlinked
   representation is a positive Float64 value (e.g. `2.0`).
2. The unlinked vectorised representation is a one-element vector containing that positive
   Float64 value (e.g. `[2.0]`).
3. For `LogNormal()`, the appropriate link function is the logarithm, which maps positive
   values to the entire real line. Thus, the linked representation is `log(2.0) ≈ 0.693`.
4. The linked vectorised representation is a one-element vector containing that real value
    (e.g. `[0.693]`).

## Example 2

Consider now `y ~ LKJCholesky(2, 0.5)`.

1. A sample from this distribution is a `LinearAlgebra.Cholesky` object which wraps a
   2x2 lower-triangular matrix. The top-left element is always 1.0, and the bottom row
   must satisfy y.L[2, 1]^2 + y.L[2, 2]^2 = 1.

2. The unlinked vectorised representation is a vectorised form of this lower-triangular
   matrix. In principle, y.L[1, 1] (which is always one) and y.L[1, 2] (which is always
   zero) could be dropped to save space, but DynamicPPL's implementation retains it. This
   Thus, the unlinked vectorised representation is a 4-element vector `[1.0, 0.0, y.L[2, 1],
   y.L[2, 2]]`.

3. The link function for `LKJCholesky` maps the Cholesky sample to a vector of unconstrained
   real values. This is accomplished using Bijectors.VecCholeskyBijector. The implementation
   of this is not important here, but it is worth pointing out that there is only one free
   parameter in this 2x2 Cholesky factor (since the top row is fixed, and the two elements in
   the lower row are interdependent). Thus, the linked representation is a vector with one
   value.

4. Since the linked representation is already a vector, the linked vectorised representation
   is identical to it.

## Transform functions

This file defines functions which map between these different representations. Specifically:
                                                                          

                             linked vectorised                                     
                                                                                   
                                    ▲ │                                            
                                    │ │                                            
          from_linked_vec_transform │ │ to_linked_vec_transform                    
                                    │ │                                            
                                    │ ▼                                            
                                                  link_transform                   
                                 unlinked     ──────────────────────►      linked  
                                              ◄──────────────────────              
                                    ▲ │          invlink_transform                 
                                    │ │                                            
                 from_vec_transform │ │ to_vec_transform                           
                                    │ │                                            
                                    │ ▼                                            
                                                                                   
                             unlinked vectorised                                   

Note that all of these `..._transform` functions do not actually perform the transformation.
Instead, when called with a distribution `dist` as the only argument, they return a new
function which _then_ performs the transformation.

For example, if `dist = LogNormal()` and `x = 2.0`, then _in principle_ the following should
hold:

- `to_vec_transform(dist)(x) == [2.0]`
- `link_transform(dist)(x) == log(2.0)`
- `to_linked_vec_transform(dist)(x) == [log(2.0)]

## Implementing a new distribution

To implement the necessary transformations for a new distribution `dist`, you need to do the
following:

1. `link_transform` and `invlink_transform` should be implemented via Bijectors.jl; there is
   no need to add custom functionality for this in DynamicPPL. In there you should define `b
   = bijector(dist)`, along with `with_logabsdet_jacobian(b, x)` where `x` is in the support
   of `dist`. You also need to define `inverse(b)` and `with_logabsdet_jacobian(inverse(b),
   y)`.

2. Implement `from_vec_transform(dist)`. It turns out that although `to_vec_transform`
   exists as a function which can be overloaded, DynamicPPL does not actually ever call it,
   so you do not need to implement it.

3. Implement `from_linked_vec_transform(dist)`. There are three general cases here, and your
   distribution will most likely fall into one of them:

   - If `dist` is univariate, then the linked representation is a scalar, and the linked
     vectorised representation is a one-element vector. In this case, you can implement
     `from_linked_vec_transform(dist)` as the composition of `invlink_transform(dist)` and
     `only`.

   - If the linked representation is already a vector (which is always the case for
     multivariate distributions), then `from_linked_vec_transform(dist) =
     invlink_transform(dist)`. This follows because the linked vectorised representation
     is already the same as the linked representation. For example, this is the case for
     `LKJCholesky`.

   - If the linked representation is something else (e.g. a matrix) then you will need to
     implement a custom `from_linked_vec_transform(dist)`, which first reshapes the linked
     vectorised representation back into the linked representation, and then applies
     `invlink_transform(dist)`.

## Jacobians

The above is slightly misleading, because the actual requirement is not to implement the
transformation functions `trf(x)` for the transformation returned by `from_vec_transform`
and friends. In fact, what DynamicPPL really needs is an implementation of
`Bijectors.with_logabsdet_jacobian(trf, x)` for the transformation `trf`. This requirement
means that it is often more convenient to implement `trf` as a callable struct, since this
allows for convenient dispatch on its type.

For simple vectorisation transforms like the outputs of `from_vec_transform`, the Jacobian
term is zero, so often `Bijectors.with_logabsdet_jacobian(trf, x)` just returns `trf(x),
zero(LogProbType)`. For the results of `from_linked_vec_transform`, the Jacobian term will
be the same as that of the link transformation. In general, you can often rely on this
to be automatically computed if you define `from_linked_vec_transform` as a composition
of other transformations for which `Bijectors.with_logabsdet_jacobian` is already defined.
"""