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
penelopeysm wants to merge 3 commits into from
Closed

Specialise univariate transforms #1145

Show file tree
Hide file tree
Changes from 1 commit
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
27 changes: 15 additions & 12 deletions src/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,18 @@ 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)

Expand All @@ -363,7 +375,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 @@ -445,18 +457,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