diff --git a/.github/workflows/FloatTypes.yml b/.github/workflows/FloatTypes.yml new file mode 100644 index 000000000..e716902e0 --- /dev/null +++ b/.github/workflows/FloatTypes.yml @@ -0,0 +1,42 @@ +name: Float type promotion + +on: + push: + branches: + - main + pull_request: + +# needed to allow julia-actions/cache to delete old caches that it has created +permissions: + actions: write + contents: read + +# Cancel existing tests on the same PR if a new commit is added to a pull request +concurrency: + group: ${{ github.workflow }}-${{ github.ref || github.run_id }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + +jobs: + floattypes: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - uses: julia-actions/setup-julia@v2 + with: + version: "1" + + - uses: julia-actions/cache@v2 + + - name: Run float type tests + working-directory: test/floattypes + run: | + julia --project=. --color=yes -e 'using Pkg; Pkg.instantiate()' + julia --project=. --color=yes main.jl setup f64 + julia --project=. --color=yes main.jl run f64 + julia --project=. --color=yes main.jl setup f32 + julia --project=. --color=yes main.jl run f32 + julia --project=. --color=yes main.jl setup f16 + julia --project=. --color=yes main.jl run f16 + julia --project=. --color=yes main.jl setup min + julia --project=. --color=yes main.jl run min diff --git a/.gitignore b/.gitignore index d5a87f1eb..10d8a5f69 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ Manifest.toml **.~undo-tree~ benchmarks/*.json +LocalPreferences.toml diff --git a/HISTORY.md b/HISTORY.md index 17c328988..5b4c37b30 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,17 @@ +# 0.40.15 + +DynamicPPL now allows you to set the type that log-probabilities are initialised with, using the `DynamicPPL.set_logprob_type!` function. +This records a compile-time preference so requires restarting Julia to take effect. + +This allows model log-probability accumulation to work with different numerical precisions. +For example, if your model is defined using distributions that are parameterised by `Float32` only (and avoid promoting them to `Float64` elsewhere in the model), and you call `DynamicPPL.set_logprob_type!(Float32)`, the resulting log-probabilities will also be `Float32`. + +Previously, DynamicPPL would automatically choose a `Float64` log-probability, causing any lower-precision model to be promoted. + +The function `DynamicPPL.get_input_vector_type(::LogDensityFunction)` is now exported, in order to help with querying the type that log-probabilities are initialised with. + +`DynamicPPL.typed_identity` is deprecated; please use `Bijectors.VectorBijectors.TypedIdentity()` instead (it does the same thing). + # 0.40.14 Fixed `check_model()` erroneously failing for models such as `x[1:2] .~ univariate_dist`. diff --git a/Project.toml b/Project.toml index af9ae267f..c794203dd 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DynamicPPL" uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8" -version = "0.40.14" +version = "0.40.15" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -22,6 +22,7 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -70,6 +71,7 @@ MarginalLogDensities = "0.4.3" Mooncake = "0.4.147, 0.5" OrderedCollections = "1" PrecompileTools = "1.2.1" +Preferences = "1.5.2" Printf = "1.10" Random = "1.6" ReverseDiff = "1" diff --git a/docs/src/api.md b/docs/src/api.md index 77066b89c..f6e66b011 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -72,6 +72,7 @@ The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) inte ```@docs LogDensityFunction +get_input_vector_type ``` Internally, this is accomplished using [`init!!`](@ref) on: @@ -193,12 +194,6 @@ DynamicPPL.prefix ## Utilities -`typed_identity` is the same as `identity`, but with an overload for `with_logabsdet_jacobian` that ensures that it never errors. - -```@docs -typed_identity -``` - It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function. ```@docs @@ -396,6 +391,14 @@ DynamicPPL.split DynamicPPL.combine ``` +The float type used for accumulation of log-probabilities is defined by a compile-time preference: + +```@docs +DynamicPPL.LogProbType +DynamicPPL.set_logprob_type! +DynamicPPL.NoLogProb +``` + ```@docs VNTAccumulator DoNotAccumulate diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index 619e8d4c2..7d26e073a 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.jl @@ -114,7 +114,6 @@ export AbstractVarInfo, @model, # Utilities OrderedDict, - typed_identity, # Model Model, getmissings, @@ -129,6 +128,7 @@ export AbstractVarInfo, LogDensityFunction, OnlyAccsVarInfo, to_vector_params, + get_input_vector_type, # Leaf contexts AbstractContext, contextualize, @@ -201,8 +201,10 @@ export AbstractVarInfo, # Convenience macros @addlogprob!, check_model, + set_logprob_type!, # Deprecated. - generated_quantities + generated_quantities, + typed_identity # Reexport using Distributions: loglikelihood diff --git a/src/accumulators/default.jl b/src/accumulators/default.jl index 0cf5dd7a7..589823bf2 100644 --- a/src/accumulators/default.jl +++ b/src/accumulators/default.jl @@ -17,13 +17,14 @@ types like LogPriorAccumulator, LogJacobianAccumulator, and LogLikelihoodAccumul """ abstract type LogProbAccumulator{T<:Real} <: AbstractAccumulator end -# The first of the below methods sets AccType{T}() = AccType(zero(T)) for any -# AccType <: LogProbAccumulator{T}. The second one sets LogProbType as the default eltype T -# when calling AccType(). """ LogProbAccumulator{T}() -Create a new `LogProbAccumulator` accumulator with the log prior initialized to zero. +Create a new `LogProbAccumulator` accumulator with the log prior initialized to `zero(T)`. + + LogProbAccumulator() + +Create a new `LogProbAccumulator{DynamicPPL.LogProbType}` accumulator. """ (::Type{AccType})() where {T<:Real,AccType<:LogProbAccumulator{T}} = AccType(zero(T)) (::Type{AccType})() where {AccType<:LogProbAccumulator} = AccType{LogProbType}() @@ -59,6 +60,7 @@ function combine(acc::LogProbAccumulator, acc2::LogProbAccumulator) end acclogp(acc::LogProbAccumulator, val) = basetypeof(acc)(logp(acc) + val) +acclogp(acc::LogProbAccumulator{NoLogProb}, val) = basetypeof(acc)(val) function Base.convert( ::Type{AccType}, acc::LogProbAccumulator diff --git a/src/accumulators/vector_params.jl b/src/accumulators/vector_params.jl index 0f7b9dd7d..2ab48816e 100644 --- a/src/accumulators/vector_params.jl +++ b/src/accumulators/vector_params.jl @@ -17,7 +17,7 @@ initialisation strategy and collect the vectorised parameters corresponding to t strategy. """ function VectorParamAccumulator(ldf::LogDensityFunction) - et = eltype(_get_input_vector_type(ldf)) + et = eltype(get_input_vector_type(ldf)) dim = ldf._dim vals = Vector{et}(undef, dim) set_indices = falses(dim) diff --git a/src/deprecated.jl b/src/deprecated.jl index 0bcaae9b7..bc978faf1 100644 --- a/src/deprecated.jl +++ b/src/deprecated.jl @@ -1 +1,3 @@ @deprecate generated_quantities(model, params) returned(model, params) + +typed_identity = Bijectors.VectorBijectors.TypedIdentity() diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index 5ff30f80f..b1ae82f2b 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -309,7 +309,17 @@ function _default_vnt(model::Model, transform_strategy::AbstractTransformStrateg return getacc(oavi, Val(VECTORVAL_ACCNAME)).values end -function _get_input_vector_type(::LogDensityFunction{M,A,L,G,R,P,X}) where {M,A,L,G,R,P,X} +""" + DynamicPPL.get_input_vector_type(::LogDensityFunction) + +Get the type of the vector `x` that should be passed to `LogDensityProblems.logdensity(ldf, +x)`. + +Note that if you pass a vector of a different type, it will be converted to the correct +type. This allows you however to determine upfront what kind of vector should be passed in. +It is also useful for determining e.g. whether Float32 or Float64 parameters are expected. +""" +function get_input_vector_type(::LogDensityFunction{M,A,L,G,R,P,X}) where {M,A,L,G,R,P,X} return X end @@ -415,7 +425,7 @@ function LogDensityProblems.logdensity_and_gradient( ) # `params` has to be converted to the same vector type that was used for AD preparation, # otherwise the preparation will not be valid. - params = convert(_get_input_vector_type(ldf), params) + params = convert(get_input_vector_type(ldf), params) return if _use_closure(ldf.adtype) DI.value_and_gradient( LogDensityAt( @@ -586,7 +596,7 @@ that. """ function to_vector_params(vector_values::VarNamedTuple, ldf::LogDensityFunction) return to_vector_params_inner( - vector_values, ldf._varname_ranges, eltype(_get_input_vector_type(ldf)), ldf._dim + vector_values, ldf._varname_ranges, eltype(get_input_vector_type(ldf)), ldf._dim ) end diff --git a/src/transformed_values.jl b/src/transformed_values.jl index 1befbf435..10d474de1 100644 --- a/src/transformed_values.jl +++ b/src/transformed_values.jl @@ -179,7 +179,7 @@ struct UntransformedValue{V} <: AbstractTransformedValue end Base.:(==)(tv1::UntransformedValue, tv2::UntransformedValue) = tv1.val == tv2.val Base.isequal(tv1::UntransformedValue, tv2::UntransformedValue) = isequal(tv1.val, tv2.val) -get_transform(::UntransformedValue) = typed_identity +get_transform(::UntransformedValue) = Bijectors.VectorBijectors.TypedIdentity() get_internal_value(tv::UntransformedValue) = tv.val set_internal_value(::UntransformedValue, new_val) = UntransformedValue(new_val) diff --git a/src/utils.jl b/src/utils.jl index f89071349..d218a9be9 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -2,48 +2,106 @@ # defined in other files. function subset end +using Preferences: @load_preference, @set_preferences! + """ -The type for all log probability variables. + DynamicPPL.NoLogProb <: Real + +Singleton type that represents the absence of a log probability value. This is used as the +default type parameter for `LogProbAccumulator` when no log probability value is needed, to +avoid defining a concrete type such as `Float64` that would cause unwanted type promotion +when accumulating log probabilities of other types (e.g., `Float32`). -This is Float64 on 64-bit systems and Float32 on 32-bit systems. +Adding anything to `NoLogProb()` returns the other thing. In other words, `NoLogProb` is a +true additive identity which additionally preserves types. """ -const LogProbType = float(Real) +struct NoLogProb <: Real end +Base.zero(::Type{NoLogProb}) = NoLogProb() +Base.convert(::Type{T}, ::NoLogProb) where {T<:Number} = zero(T) +Base.promote_rule(::Type{NoLogProb}, ::Type{T}) where {T<:Number} = T +Base.iszero(::NoLogProb) = true +Base.hash(::NoLogProb, h::UInt) = hash(0.0, h) +Base.:(+)(::NoLogProb, ::NoLogProb) = NoLogProb() +Base.:(-)(::NoLogProb, ::NoLogProb) = NoLogProb() + +const FLOAT_TYPE_PREF_KEY = "floattype" """ - typed_identity(x) + DynamicPPL.LogProbType -Identity function, but with an overload for `with_logabsdet_jacobian` to ensure -that it returns a sensible zero logjac. +The default type used for log-probabilities in DynamicPPL.jl. This is a compile-time constant +that can be set via [`set_logprob_type!`](@ref), which under the hood uses Preferences.jl. -The problem with plain old `identity` is that the default definition of -`with_logabsdet_jacobian` for `identity` returns `zero(eltype(x))`: -https://github.com/JuliaMath/ChangesOfVariables.jl/blob/d6a8115fc9b9419decbdb48e2c56ec9675b4c6a4/src/with_ladj.jl#L154 +Note that this does not prevent computations within the model from promoting the +log-probability to a different type. In essence, `LogProbType` specifies the *lowest* +possible type that log-probabilities can be, and DynamicPPL promises to not insert any extra +operations that would cause this to be promoted to a higher type. However, DynamicPPL cannot +guard against user code inside models. -This is fine for most samples `x`, but if `eltype(x)` doesn't return a sensible type (e.g. -if it's `Any`), then using `identity` will error with `zero(Any)`. This can happen with, -for example, `ProductNamedTupleDistribution`: +For example, in: ```julia -julia> using Distributions; d = product_distribution((a = Normal(), b = LKJCholesky(3, 0.5))); +@model f() = x ~ Normal(0.0, 1.0) +``` + +the log-probability of the model will always be promoted to `Float64`, regardless of the +value of `LogProbType`, because the logpdf of `Normal(0.0, 1.0)` is a `Float64`. On the +other hand, in: -julia> eltype(rand(d)) -Any +```julia +@model f() = x ~ Normal(0.0f0, 1.0f0) ``` -The same problem precludes us from eventually broadening the scope of DynamicPPL.jl to -support distributions with non-numeric samples. +the log-probability of the model will be `Float32` if `LogProbType` is `Float32` or lower. +""" +const LogProbType = let + logp_pref = @load_preference(FLOAT_TYPE_PREF_KEY, "f64") + if logp_pref == "f64" + Float64 + elseif logp_pref == "f32" + Float32 + elseif logp_pref == "f16" + Float16 + elseif logp_pref == "min" + NoLogProb + else + error("Unsupported log probability preference: $logp_pref") + end +end -Furthermore, in principle, the type of the log-probability should be separate from the type -of the sample. Thus, instead of using `zero(LogProbType)`, we should use the eltype of the -LogJacobianAccumulator. There's no easy way to thread that through here, but if a way to do -this is discovered, then `typed_identity` is what will allow us to obtain that custom -behaviour. """ -function typed_identity end -@inline typed_identity(x) = x -@inline Bijectors.with_logabsdet_jacobian(::typeof(typed_identity), x) = - (x, zero(LogProbType)) -@inline Bijectors.inverse(::typeof(typed_identity)) = typed_identity + set_logprob_type!(::Type{T}) where {T} + +Set the log probability type for DynamicPPL.jl, [`DynamicPPL.LogProbType`](@ref), to `T`. +Permitted values are `Float64`, `Float32`, `Float16`, and `NoLogProb`. The default in +DynamicPPL is `Float64`. + +`NoLogProb` is a special type that is the "lowest" possible float type. This means that the +log probability will be promoted to whatever type the model dictates. This is a totally +unintrusive option, which can be useful if you do not know in advance what log probability +type you are targeting, or want to troubleshoot a model to see what type the log probability +is being promoted to. However, this can also cause type stability issues and performance +degradations, so we generally recommend setting a specific log probability type if you know +what type you want to target. + +This function uses Preferences.jl to set a compile-time constant, so you will need to +restart your Julia session for the change to take effect. +""" +function set_logprob_type!(::Type{T}) where {T} + new_pref = if T == Float64 + "f64" + elseif T == Float32 + "f32" + elseif T == Float16 + "f16" + elseif T == NoLogProb + "min" + else + throw(ArgumentError("Unsupported log probability type: $T")) + end + @set_preferences!(FLOAT_TYPE_PREF_KEY => new_pref) + @info "DynamicPPL's log probability type has been set to $T.\nPlease note you will need to restart your Julia session for this change to take effect." +end """ @addlogprob!(ex) diff --git a/test/accumulators.jl b/test/accumulators.jl index b88d76e96..6ecde3e84 100644 --- a/test/accumulators.jl +++ b/test/accumulators.jl @@ -11,6 +11,7 @@ using DynamicPPL: AccumulatorTuple, LogLikelihoodAccumulator, LogPriorAccumulator, + NoLogProb, accumulate_assume!!, accumulate_observe!!, combine, @@ -39,6 +40,25 @@ using DynamicPPL: DynamicPPL.reset(LogLikelihoodAccumulator(1.0)) end + @testset "float types" begin + # f64 + @test DynamicPPL.reset(LogPriorAccumulator(1.0)) === LogPriorAccumulator(0.0) + @test DynamicPPL.reset(LogPriorAccumulator()) === LogPriorAccumulator() + + # f32 + @test LogPriorAccumulator{Float32}() == LogPriorAccumulator(0.0f0) + @test DynamicPPL.reset(LogPriorAccumulator(1.0f0)) === + LogPriorAccumulator(0.0f0) + @test DynamicPPL.reset(LogPriorAccumulator{Float32}()) === + LogPriorAccumulator{Float32}() + + # nologprob + @test LogPriorAccumulator(NoLogProb()) !== LogPriorAccumulator(0.0) + @test LogPriorAccumulator(NoLogProb()) === LogPriorAccumulator(NoLogProb()) + @test DynamicPPL.reset(LogPriorAccumulator(NoLogProb())) === + LogPriorAccumulator(NoLogProb()) + end + @testset "addition and incrementation" begin @test acclogp(LogPriorAccumulator(1.0f0), 1.0f0) == LogPriorAccumulator(2.0f0) @test acclogp(LogPriorAccumulator(1.0), 1.0f0) == LogPriorAccumulator(2.0) @@ -48,6 +68,13 @@ using DynamicPPL: LogLikelihoodAccumulator(2.0) end + @testset "addition to NoLogProb" begin + for val in (1.0, 1.0f0, BigFloat(1.0), 1, UInt(1), Float16(1.0)) + @test acclogp(LogPriorAccumulator(NoLogProb()), val) == + LogPriorAccumulator(val) + end + end + @testset "split and combine" begin for acc in [ LogPriorAccumulator(1.0), diff --git a/test/floattypes/Project.toml b/test/floattypes/Project.toml new file mode 100644 index 000000000..02a770fe7 --- /dev/null +++ b/test/floattypes/Project.toml @@ -0,0 +1,10 @@ +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +DynamicPPL = {path = "../../"} diff --git a/test/floattypes/main.jl b/test/floattypes/main.jl new file mode 100644 index 000000000..235bf5b52 --- /dev/null +++ b/test/floattypes/main.jl @@ -0,0 +1,79 @@ +# Script to use for testing promotion of log-prob types. Since this relies on compile-time +# preferences, it's hard to run this within the usual CI setup. +# +# Usage: +# julia --project=. main.jl setup f32 # Sets the preference +# julia --project=. main.jl run f32 # Checks that the preference is respected +# +# and this should be looped over for `f64`, `f32`, `f16`, and `min`. + +using DynamicPPL, LogDensityProblems, ForwardDiff, Distributions, ADTypes, Test + +function floattypestr_to_type(floattypestr) + if floattypestr == "f64" + return Float64 + elseif floattypestr == "f32" + return Float32 + elseif floattypestr == "f16" + return Float16 + elseif floattypestr == "min" + return DynamicPPL.NoLogProb + else + error("Invalid float type: $floattypestr") + end +end + +function setup(floattypestr) + T = floattypestr_to_type(floattypestr) + return DynamicPPL.set_logprob_type!(T) +end + +function test_with_type(::Type{T}) where {T} + @testset "Testing with type $T" begin + @model f() = x ~ Normal(T(0), T(1)) + model = f() + vnt = rand(model) + @test vnt[@varname(x)] isa T + lj = (@inferred logjoint(f(), (; x=T(0.0)))) + @test lj isa T + ldf = LogDensityFunction( + f(), getlogjoint_internal, LinkAll(); adtype=AutoForwardDiff() + ) + @test rand(ldf) isa AbstractVector{T} + lp = (@inferred LogDensityProblems.logdensity(ldf, [T(0)])) + @test lp isa T + @test lp ≈ logpdf(Normal(T(0), T(1)), T(0)) + lp_and_grad = (@inferred LogDensityProblems.logdensity_and_gradient(ldf, [T(0)])) + @test first(lp_and_grad) isa T + @test last(lp_and_grad) isa AbstractVector{T} + end +end + +function run(floattypestr) + T = floattypestr_to_type(floattypestr) + if T == DynamicPPL.NoLogProb + @test DynamicPPL.LogProbType === DynamicPPL.NoLogProb + # all higher types should cause promotion to those types + test_with_type(Float16) + test_with_type(Float32) + test_with_type(Float64) + else + @test DynamicPPL.LogProbType === T + test_with_type(T) + end +end + +if length(ARGS) != 2 || + !(ARGS[1] in ["setup", "run"]) || + !(ARGS[2] in ["f64", "f32", "f16", "min"]) + println("Usage: julia --project=. main.jl ") + exit(1) +end + +mode = ARGS[1] +floattypestr = ARGS[2] +if mode == "setup" + setup(floattypestr) +elseif mode == "run" + run(floattypestr) +end diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 904c2d69d..7ccf4efcb 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -96,6 +96,7 @@ end ldf = DynamicPPL.LogDensityFunction(model) xs = [1.0] + @test LogDensityProblems.logdensity(ldf, xs) ≈ logpdf(Normal(), xs[1]) + N * logpdf(Normal(xs[1]), 0.0) end