Skip to content

Specialise univariate transforms #1145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

Specialise univariate transforms #1145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
866a7d0
optimise univariate transforms
penelopeysm Nov 18, 2025
3573a0e
input length for only
penelopeysm Nov 18, 2025
a302b9d
remove UnwrapSingletonTransform
penelopeysm Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions HISTORY.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
# DynamicPPL Changelog

## 0.38.10

Improved performance of transformations of univariate distributions' samples to and from their vectorised forms.

## 0.38.9

Remove warning when using Enzyme as the AD backend.
Expand Down
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "DynamicPPL"
uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
version = "0.38.9"
version = "0.38.10"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand Down
71 changes: 19 additions & 52 deletions src/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@ const NO_DEFAULT = NoDefault()
VarNameTuple = NTuple{N,VarName} where {N}

# TODO(mhauru) This is currently used in the transformation functions of NoDist,
# ReshapeTransform, and UnwrapSingletonTransform, and in VarInfo. We should also use it in
# SimpleVarInfo and maybe other places.
# ReshapeTransform, and in VarInfo. We should also use it in SimpleVarInfo and maybe other
# places.
"""
The type for all log probability variables.

Expand Down Expand Up @@ -229,42 +229,6 @@ invlink_transform(dist) = inverse(link_transform(dist))
# Helper functions for vectorize/reconstruct values #
#####################################################

"""
UnwrapSingletonTransform(input_size::InSize)

A transformation that unwraps a singleton array, returning a scalar.

The `input_size` field is the expected size of the input. In practice this only determines
the number of indices, since all dimensions must be 1 for a singleton. `input_size` is used
to check the validity of the input, but also to determine the correct inverse operation.

By default `input_size` is `(1,)`, in which case `tovec` is the inverse.
"""
struct UnwrapSingletonTransform{InSize} <: Bijectors.Bijector
input_size::InSize
end

UnwrapSingletonTransform() = UnwrapSingletonTransform((1,))

function (f::UnwrapSingletonTransform)(x)
if size(x) != f.input_size
throw(DimensionMismatch("Expected input of size $(f.input_size), got $(size(x))"))
end
return only(x)
end

function Bijectors.with_logabsdet_jacobian(f::UnwrapSingletonTransform, x)
return f(x), zero(LogProbType)
end

function Bijectors.with_logabsdet_jacobian(
inv_f::Bijectors.Inverse{<:UnwrapSingletonTransform}, x
)
f = inv_f.orig
result = reshape([x], f.input_size)
return result, zero(LogProbType)
end

"""
ReshapeTransform(input_size::InSize, output_size::OutSize)

Expand Down Expand Up @@ -338,14 +302,26 @@ function Bijectors.with_logabsdet_jacobian(::Bijectors.Inverse{<:ToChol}, y)
)
end

struct Only end
struct NotOnly end
(::Only)(x) = x[]
(::NotOnly)(y) = [y]
function Bijectors.with_logabsdet_jacobian(::Only, x::AbstractVector{T}) where {T<:Real}
return (x[], zero(T))
end
Bijectors.with_logabsdet_jacobian(::Only, x::AbstractVector) = (x[], zero(LogProbType))
Bijectors.inverse(::Only) = NotOnly()
Bijectors.with_logabsdet_jacobian(::NotOnly, y::T) where {T<:Real} = ([y], zero(T))
Bijectors.with_logabsdet_jacobian(::NotOnly, y) = ([y], zero(LogProbType))

"""
from_vec_transform(x)

Return the transformation from the vector representation of `x` to original representation.
"""
from_vec_transform(x::AbstractArray) = from_vec_transform_for_size(size(x))
from_vec_transform(C::Cholesky) = ToChol(C.uplo) ∘ ReshapeTransform(size(C.UL))
from_vec_transform(::Real) = UnwrapSingletonTransform()
from_vec_transform(::Real) = Only()

"""
from_vec_transform_for_size(sz::Tuple)
Expand All @@ -363,7 +339,7 @@ Return the transformation from the vector representation of a realization from
distribution `dist` to the original representation compatible with `dist`.
"""
from_vec_transform(dist::Distribution) = from_vec_transform_for_size(size(dist))
from_vec_transform(::UnivariateDistribution) = UnwrapSingletonTransform()
from_vec_transform(::UnivariateDistribution) = Only()
from_vec_transform(dist::LKJCholesky) = ToChol(dist.uplo) ∘ ReshapeTransform(size(dist))

struct ProductNamedTupleUnvecTransform{names,T<:NamedTuple{names}}
Expand Down Expand Up @@ -409,7 +385,7 @@ end
# This function returns the length of the vector that the function from_vec_transform
# expects. This helps us determine which segment of a concatenated vector belongs to which
# variable.
_input_length(from_vec_trfm::UnwrapSingletonTransform) = 1
_input_length(::Only) = 1
_input_length(from_vec_trfm::ReshapeTransform) = prod(from_vec_trfm.output_size)
function _input_length(trfm::ProductNamedTupleUnvecTransform)
return sum(_input_length ∘ from_vec_transform, values(trfm.dists))
Expand Down Expand Up @@ -445,18 +421,9 @@ function from_linked_vec_transform(dist::Distribution)
f_vec = from_vec_transform(inverse(f_invlink), size(dist))
return f_invlink ∘ f_vec
end

# UnivariateDistributions need to be handled as a special case, because size(dist) is (),
# which makes the usual machinery think we are dealing with a 0-dim array, whereas in
# actuality we are dealing with a scalar.
Comment on lines -449 to -451
Copy link
Member Author

@penelopeysm penelopeysm Nov 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i mean isn't this just annoying, this is what i mean when i say bijectors is a pain to deal with

Copy link
Member Author

@penelopeysm penelopeysm Nov 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and also distributions to a lesser extent

size(dist) and Bijectors.output_size are not well-defined for things like LKJCholesky or ProductNamedTupleDistribution

that's fine because their use cases are (possibly) different, but dynamicppl has its own needs, and the correct solution is to have a clean api that generalises across all of these cases

# TODO(mhauru) Hopefully all this can go once the old Gibbs sampler is removed and
# VarNamedVector takes over from Metadata.
function from_linked_vec_transform(dist::UnivariateDistribution)
f_invlink = invlink_transform(dist)
f_vec = from_vec_transform(inverse(f_invlink), size(dist))
f_combined = f_invlink ∘ f_vec
sz = Bijectors.output_size(f_combined, size(dist))
return UnwrapSingletonTransform(sz) ∘ f_combined
# This is a performance optimisation
return Only() ∘ invlink_transform(dist)
end
function from_linked_vec_transform(dist::Distributions.ProductNamedTupleDistribution)
return invlink_transform(dist)
Expand Down
Loading