diff --git a/.github/workflows/FloatTypes.yml b/.github/workflows/FloatTypes.yml new file mode 100644 index 0000000000..e716902e04 --- /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 23d6e12f37..515d452704 100644 --- a/.gitignore +++ b/.gitignore @@ -25,7 +25,6 @@ docs/.jekyll-cache .vscode .DS_Store Manifest.toml -/Manifest.toml -/test/Manifest.toml +LocalPreferences.toml benchmarks/output/ diff --git a/HISTORY.md b/HISTORY.md index a6b5e3f1fe..15d0f7fb62 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,15 @@ +# 0.43.3 + +Unify parameter initialisation for HMC and external samplers. +External samplers (like HMC) now attempt multiple times to generate valid initial parameters, instead of just taking the first set of parameters. + +Re-exports `set_logprob_type!` from DynamicPPL to allow users to control the base log-probability type used when evaluating Turing models. +For example, calling `set_logprob_type!(Float32)` will mean that Turing will use `Float32` for log-probability calculations, only promoting if there is something in the model that causes it to be (e.g. a distribution that returns `Float64` log-probabilities). +Note that this is a compile-time preference: for it to take effect you will have to restart your Julia session after calling `set_logprob_type!`. + +Furthermore, note that sampler support for non-`Float64` log-probabilities is currently limited. +Although DynamicPPL promises not promote float types unnecessarily, many samplers, including HMC and NUTS, still use `Float64` internally and thus will cause log-probabilities and parameters to be promoted to `Float64`, even if the model itself uses `Float32`. + # 0.43.2 Throw an `ArgumentError` when a `Gibbs` sampler is missing component samplers for any variable in the model. diff --git a/Project.toml b/Project.toml index f9d396cc11..69b242a6f8 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.43.2" +version = "0.43.3" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -61,7 +61,7 @@ DifferentiationInterface = "0.7" Distributions = "0.25.77" DocStringExtensions = "0.8, 0.9" DynamicHMC = "3.4" -DynamicPPL = "0.40.6" +DynamicPPL = "0.40.15" EllipticalSliceSampling = "0.5, 1, 2" ForwardDiff = "0.10.3, 1" Libtask = "0.9.14" diff --git a/docs/src/api.md b/docs/src/api.md index ca08e166bb..b92a0b7acd 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -46,6 +46,7 @@ even though [`Prior()`](@ref) is actually defined in the `Turing.Inference` modu | `setthreadsafe` | [`DynamicPPL.setthreadsafe`](@extref) | Mark a model as requiring threadsafe evaluation | | `might_produce` | [`Libtask.might_produce`](@extref) | Mark a method signature as potentially calling `Libtask.produce` | | `@might_produce` | [`Libtask.@might_produce`](@extref) | Mark a function name as potentially calling `Libtask.produce` | +| `set_logprob_type!` | [`DynamicPPL.set_logprob_type!`](@extref) | Set the base log-probability type used during evaluation of Turing models | ### Inference @@ -81,11 +82,11 @@ even though [`Prior()`](@ref) is actually defined in the `Turing.Inference` modu ### Data structures -| Exported symbol | Documentation | Description | -|:--------------- |:------------------------------------------- |:----------------------------------- | -| `@vnt` | [`DynamicPPL.@vnt`](@extref) | Generate a `VarNameTuple` | -| `VarNamedTuple` | [`DynamicPPL.VarNamedTuple`](@extref) | A mapping from `VarName`s to values | -| `OrderedDict` | [`OrderedCollections.OrderedDict`](@extref) | An ordered dictionary | +| Exported symbol | Documentation | Description | +|:--------------- |:---------------------------------------------------- |:----------------------------------- | +| `@vnt` | [`DynamicPPL.VarNamedTuples.@vnt`](@extref) | Generate a `VarNameTuple` | +| `VarNamedTuple` | [`DynamicPPL.VarNamedTuples.VarNamedTuple`](@extref) | A mapping from `VarName`s to values | +| `OrderedDict` | [`OrderedCollections.OrderedDict`](@extref) | An ordered dictionary | ### DynamicPPL utilities diff --git a/ext/TuringDynamicHMCExt.jl b/ext/TuringDynamicHMCExt.jl index 041aa0865d..b4cfadaf61 100644 --- a/ext/TuringDynamicHMCExt.jl +++ b/ext/TuringDynamicHMCExt.jl @@ -50,26 +50,23 @@ function AbstractMCMC.step( initial_params, kwargs..., ) - # Define log-density function. - # TODO(penelopeysm) We need to check that the initial parameters are valid. Same as how - # we do it for HMC - _, vi = DynamicPPL.init!!( - rng, model, DynamicPPL.VarInfo(), initial_params, DynamicPPL.LinkAll() - ) - ℓ = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype + # Construct LogDensityFunction + tfm_strategy = DynamicPPL.LinkAll() + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, tfm_strategy; adtype=spl.adtype ) + x = Turing.Inference.find_initial_params_ldf(rng, ldf, initial_params) # Perform initial step. results = DynamicHMC.mcmc_keep_warmup( - rng, ℓ, 0; initialization=(q=vi[:],), reporter=DynamicHMC.NoProgressReport() + rng, ldf, 0; initialization=(q=x,), reporter=DynamicHMC.NoProgressReport() ) steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state) Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q) # Create first sample and state. - sample = DynamicPPL.ParamsWithStats(Q.q, ℓ) - state = DynamicNUTSState(ℓ, Q, steps.H.κ, steps.ϵ) + sample = DynamicPPL.ParamsWithStats(Q.q, ldf) + state = DynamicNUTSState(ldf, Q, steps.H.κ, steps.ϵ) return sample, state end diff --git a/src/Turing.jl b/src/Turing.jl index 87f09a4533..0b15e80004 100644 --- a/src/Turing.jl +++ b/src/Turing.jl @@ -78,7 +78,9 @@ using DynamicPPL: InitFromParams, setthreadsafe, filldist, - arraydist + arraydist, + set_logprob_type! + using StatsBase: predict using OrderedCollections: OrderedDict using Libtask: might_produce, @might_produce @@ -163,6 +165,8 @@ export fix, unfix, OrderedDict, # OrderedCollections + # Log-prob types in accumulators + set_logprob_type!, # Initialisation strategies for models InitFromPrior, InitFromUniform, diff --git a/src/mcmc/abstractmcmc.jl b/src/mcmc/abstractmcmc.jl index f867ab5344..ac99ec3ef2 100644 --- a/src/mcmc/abstractmcmc.jl +++ b/src/mcmc/abstractmcmc.jl @@ -31,6 +31,44 @@ function _convert_initial_params(@nospecialize(_::Any)) throw(ArgumentError(errmsg)) end +""" + find_initial_params_ldf(rng, ldf, init_strategy; max_attempts=1000) + +Given a `LogDensityFunction` and an initialization strategy, attempt to find valid initial +parameters by sampling from the initialization strategy and checking that the log density +(and gradient, if available) are finite. If valid parameters are not found after +`max_attempts`, throw an error. +""" +function find_initial_params_ldf( + rng::Random.AbstractRNG, + ldf::DynamicPPL.LogDensityFunction, + init_strategy::DynamicPPL.AbstractInitStrategy; + max_attempts::Int=1000, +) + for attempts in 1:max_attempts + # Get new parameters + x = rand(rng, ldf, init_strategy) + is_valid = if ldf.adtype === nothing + logp = LogDensityProblems.logdensity(ldf, x) + isfinite(logp) + else + logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, x) + isfinite(logp) && all(isfinite, grad) + end + + # If they're OK, return them + is_valid && return x + + attempts == 10 && + @warn "failed to find valid initial parameters in $(attempts) tries; consider providing a different initialisation strategy with the `initial_params` keyword" + end + + # if we failed to find valid initial parameters, error + return error( + "failed to find valid initial parameters in $(max_attempts) tries. See https://turinglang.org/docs/uri/initial-parameters for common causes and solutions. If the issue persists, please open an issue at https://github.com/TuringLang/Turing.jl/issues", + ) +end + ######################################### # Default definitions for the interface # ######################################### diff --git a/src/mcmc/external_sampler.jl b/src/mcmc/external_sampler.jl index 2037260e0f..f654d1ea54 100644 --- a/src/mcmc/external_sampler.jl +++ b/src/mcmc/external_sampler.jl @@ -149,28 +149,17 @@ function AbstractMCMC.step( ) where {unconstrained} sampler = sampler_wrapper.sampler - # Initialise varinfo with initial params and link the varinfo if needed. - tfm_strategy = unconstrained ? DynamicPPL.LinkAll() : DynamicPPL.UnlinkAll() - _, varinfo = DynamicPPL.init!!(rng, model, VarInfo(), initial_params, tfm_strategy) - - # We need to extract the vectorised initial_params, because the later call to - # AbstractMCMC.step only sees a `LogDensityModel` which expects `initial_params` - # to be a vector. - initial_params_vector = varinfo[:] - # Construct LogDensityFunction + tfm_strategy = unconstrained ? DynamicPPL.LinkAll() : DynamicPPL.UnlinkAll() f = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, varinfo; adtype=sampler_wrapper.adtype + model, DynamicPPL.getlogjoint_internal, tfm_strategy; adtype=sampler_wrapper.adtype ) + x = find_initial_params_ldf(rng, f, initial_params) # Then just call `AbstractMCMC.step` with the right arguments. _, state_inner = if initial_state === nothing AbstractMCMC.step( - rng, - AbstractMCMC.LogDensityModel(f), - sampler; - initial_params=initial_params_vector, - kwargs..., + rng, AbstractMCMC.LogDensityModel(f), sampler; initial_params=x, kwargs... ) else @@ -179,7 +168,7 @@ function AbstractMCMC.step( AbstractMCMC.LogDensityModel(f), sampler, initial_state; - initial_params=initial_params_vector, + initial_params=x, kwargs..., ) end @@ -191,7 +180,13 @@ function AbstractMCMC.step( new_stats = AbstractMCMC.getstats(state_inner) DynamicPPL.ParamsWithStats(new_parameters, f, new_stats) end - return (new_transition, TuringState(state_inner, varinfo, new_parameters, f)) + + # TODO(penelopeysm): this varinfo is only needed for Gibbs. The external sampler itself + # has no use for it. Get rid of this as soon as possible. + vi = DynamicPPL.link!!(VarInfo(model), model) + vi = DynamicPPL.unflatten!!(vi, x) + + return (new_transition, TuringState(state_inner, vi, new_parameters, f)) end function AbstractMCMC.step( diff --git a/src/mcmc/hmc.jl b/src/mcmc/hmc.jl index 8be1b48784..aec5397dcc 100644 --- a/src/mcmc/hmc.jl +++ b/src/mcmc/hmc.jl @@ -150,41 +150,10 @@ function AbstractMCMC.sample( end end -function find_initial_params( - rng::Random.AbstractRNG, - model::DynamicPPL.Model, - varinfo::DynamicPPL.AbstractVarInfo, - hamiltonian::AHMC.Hamiltonian, - init_strategy::DynamicPPL.AbstractInitStrategy; - max_attempts::Int=1000, -) - varinfo = deepcopy(varinfo) # Don't mutate - - for attempts in 1:max_attempts - theta = varinfo[:] - z = AHMC.phasepoint(rng, theta, hamiltonian) - isfinite(z) && return varinfo, z - - attempts == 10 && - @warn "failed to find valid initial parameters in $(attempts) tries; consider providing a different initialisation strategy with the `initial_params` keyword" - - # Resample and try again. - _, varinfo = DynamicPPL.init!!( - rng, model, varinfo, init_strategy, DynamicPPL.LinkAll() - ) - end - - # if we failed to find valid initial parameters, error - return error( - "failed to find valid initial parameters in $(max_attempts) tries. See https://turinglang.org/docs/uri/initial-parameters for common causes and solutions. If the issue persists, please open an issue at https://github.com/TuringLang/Turing.jl/issues", - ) -end - -function Turing.Inference.initialstep( +function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, - spl::Hamiltonian, - vi_original::AbstractVarInfo; + spl::Hamiltonian; # the initial_params kwarg is always passed on from sample(), cf. DynamicPPL # src/sampler.jl, so we don't need to provide a default value here initial_params::DynamicPPL.AbstractInitStrategy, @@ -193,32 +162,19 @@ function Turing.Inference.initialstep( verbose::Bool=true, kwargs..., ) - # Transform the samples to unconstrained space and compute the joint log probability. - vi = DynamicPPL.link(vi_original, model) - - # Extract parameters. - theta = vi[:] - - # Create a Hamiltonian. - metricT = getmetricT(spl) - metric = metricT(length(theta)) + # Create a Hamiltonian ldf = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype + model, DynamicPPL.getlogjoint_internal, DynamicPPL.LinkAll(); adtype=spl.adtype ) + metricT = getmetricT(spl) + metric = metricT(LogDensityProblems.dimension(ldf)) lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf) lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf) hamiltonian = AHMC.Hamiltonian(metric, lp_func, lp_grad_func) - # Note that there is already one round of 'initialisation' before we reach this step, - # inside DynamicPPL's `AbstractMCMC.step` implementation. That leads to a possible issue - # that this `find_initial_params` function might override the parameters set by the - # user. - # Luckily for us, `find_initial_params` always checks if the logp and its gradient are - # finite. If it is already finite with the params inside the current `vi`, it doesn't - # attempt to find new ones. This means that the parameters passed to `sample()` will be - # respected instead of being overridden here. - vi, z = find_initial_params(rng, model, vi, hamiltonian, initial_params) - theta = vi[:] + # Find initial values + theta = find_initial_params_ldf(rng, ldf, initial_params) + z = AHMC.phasepoint(rng, theta, hamiltonian) # Find good eps if not provided one if iszero(spl.ϵ) @@ -236,6 +192,12 @@ function Turing.Inference.initialstep( else DynamicPPL.ParamsWithStats(theta, ldf, NamedTuple()) end + + # TODO(penelopeysm): this varinfo is only needed for Gibbs. HMC itself has no use for + # it. Get rid of this as soon as possible. + vi = DynamicPPL.link!!(VarInfo(model), model) + vi = DynamicPPL.unflatten!!(vi, theta) + state = HMCState(vi, 0, kernel, hamiltonian, z, adaptor, ldf) return transition, state diff --git a/src/mcmc/sghmc.jl b/src/mcmc/sghmc.jl index 66b6d5e5fb..35ed102daf 100644 --- a/src/mcmc/sghmc.jl +++ b/src/mcmc/sghmc.jl @@ -51,26 +51,22 @@ struct SGHMCState{L,V<:AbstractVector{<:Real},T<:AbstractVector{<:Real}} velocity::T end -function Turing.Inference.initialstep( +function AbstractMCMC.step( rng::Random.AbstractRNG, model::Model, - spl::SGHMC, - vi::AbstractVarInfo; + spl::SGHMC; + initial_params::DynamicPPL.AbstractInitStrategy, discard_sample=false, kwargs..., ) - # Transform the samples to unconstrained space. - if !DynamicPPL.is_transformed(vi) - vi = DynamicPPL.link!!(vi, model) - end - - # Compute initial sample and state. - ℓ = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype + tfm_strategy = DynamicPPL.LinkAll() + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, tfm_strategy; adtype=spl.adtype ) - initial_params = vi[:] - sample = discard_sample ? nothing : DynamicPPL.ParamsWithStats(initial_params, ℓ) - state = SGHMCState(ℓ, initial_params, zero(vi[:])) + x = Turing.Inference.find_initial_params_ldf(rng, ldf, initial_params) + + sample = discard_sample ? nothing : DynamicPPL.ParamsWithStats(x, ldf) + state = SGHMCState(ldf, x, zero(x)) return sample, state end @@ -189,31 +185,27 @@ struct SGLDState{L,V<:AbstractVector{<:Real}} step::Int end -function Turing.Inference.initialstep( +function AbstractMCMC.step( rng::Random.AbstractRNG, model::Model, - spl::SGLD, - vi::AbstractVarInfo; + spl::SGLD; + initial_params::DynamicPPL.AbstractInitStrategy, discard_sample=false, kwargs..., ) - # Transform the samples to unconstrained space. - if !DynamicPPL.is_transformed(vi) - vi = DynamicPPL.link!!(vi, model) - end - - # Create first sample and state. - ℓ = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype + tfm_strategy = DynamicPPL.LinkAll() + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, tfm_strategy; adtype=spl.adtype ) - initial_params = vi[:] + x = Turing.Inference.find_initial_params_ldf(rng, ldf, initial_params) + transition = if discard_sample nothing else stats = (; SGLD_stepsize=zero(spl.stepsize(0))) - DynamicPPL.ParamsWithStats(initial_params, ℓ, stats) + DynamicPPL.ParamsWithStats(x, ldf, stats) end - state = SGLDState(ℓ, initial_params, 1) + state = SGLDState(ldf, x, 1) return transition, state end diff --git a/src/optimisation/Optimisation.jl b/src/optimisation/Optimisation.jl index 5a925b214f..8994888d41 100644 --- a/src/optimisation/Optimisation.jl +++ b/src/optimisation/Optimisation.jl @@ -338,12 +338,7 @@ function estimate_mode( # Generate a LogDensityFunction first. We do this first because we want to use the # info stored in the LDF to generate the initial parameters and constraints in the # correct order. - vi = VarInfo(model) - vi = if link - DynamicPPL.link!!(vi, model) - else - vi - end + tfm_strategy = link ? DynamicPPL.LinkAll() : DynamicPPL.UnlinkAll() getlogdensity = logprob_func(estimator) accs = if check_constraints_at_runtime (logprob_accs(estimator)..., ConstraintCheckAccumulator(lb, ub)) @@ -352,7 +347,7 @@ function estimate_mode( end # Note that we don't need adtype to construct the LDF, because it's specified inside the # OptimizationProblem. - ldf = LogDensityFunction(model, getlogdensity, vi, accs) + ldf = LogDensityFunction(model, getlogdensity, tfm_strategy, accs) # Generate bounds and initial parameters in the unlinked or linked space as requested. lb_vec, ub_vec, inits_vec = make_optim_bounds_and_init( diff --git a/src/optimisation/init.jl b/src/optimisation/init.jl index a88ffccf56..982d3d8f51 100644 --- a/src/optimisation/init.jl +++ b/src/optimisation/init.jl @@ -317,9 +317,11 @@ function make_optim_bounds_and_init( # ranges stored in the LDF. constraint_acc = DynamicPPL.getacc(vi, Val(CONSTRAINT_ACC_NAME)) nelems = LogDensityProblems.dimension(ldf) - inits = fill(NaN, nelems) - lb = fill(-Inf, nelems) - ub = fill(Inf, nelems) + # TODO(penelopeysm) This should really be exported + et = eltype(DynamicPPL.get_input_vector_type(ldf)) + inits = fill(et(NaN), nelems) + lb = fill(et(-Inf), nelems) + ub = fill(et(Inf), nelems) for (vn, init_val) in constraint_acc.init_vecs range = _get_ldf_range(ldf, vn) inits[range] = init_val diff --git a/test/floattypes/Project.toml b/test/floattypes/Project.toml new file mode 100644 index 0000000000..eb8425d3a8 --- /dev/null +++ b/test/floattypes/Project.toml @@ -0,0 +1,7 @@ +[deps] +DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" + +[sources] +Turing = {path = "../../"} diff --git a/test/floattypes/main.jl b/test/floattypes/main.jl new file mode 100644 index 0000000000..34c0b02668 --- /dev/null +++ b/test/floattypes/main.jl @@ -0,0 +1,78 @@ +# 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, Turing, 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 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)) + + # optimisation + mr = maximum_a_posteriori(f()) + @test mr.params[@varname(x)] isa T + @test mr.lp isa T + @test eltype(mr.optim_result.u) == T + + # MH and Prior + for spl in (MH(), Prior()) + chn = sample(f(), spl, 10; progress=false) + @test eltype(chn.value) == T + end + + # nothing else really works right now + 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/main.jl b/test/main.jl new file mode 100644 index 0000000000..235bf5b522 --- /dev/null +++ b/test/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/mcmc/abstractmcmc.jl b/test/mcmc/abstractmcmc.jl index d743f25318..04478b58f2 100644 --- a/test/mcmc/abstractmcmc.jl +++ b/test/mcmc/abstractmcmc.jl @@ -2,8 +2,9 @@ module TuringAbstractMCMCTests using AbstractMCMC: AbstractMCMC using DynamicPPL: DynamicPPL -using Random: AbstractRNG -using Test: @test, @testset, @test_throws +using LogDensityProblems: LogDensityProblems +using Random: AbstractRNG, Random, Xoshiro +using Test: @test, @testset, @test_throws, @test_logs using Turing @testset "Disabling check_model" begin @@ -24,6 +25,58 @@ using Turing @test sample(model, spl, MCMCDistributed(), 10, 2; check_model=false) isa Any end +@testset "find_initial_params_ldf" begin + @testset "basic interface" begin + @model function normal_model() + x ~ Normal(0, 1) + return y ~ Normal(x, 1) + end + @testset for init_strategy in + (InitFromPrior(), InitFromUniform(), InitFromParams((x=0.5, y=-0.3))) + model = normal_model() + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, DynamicPPL.LinkAll() + ) + rng = Xoshiro(468) + x = Turing.Inference.find_initial_params_ldf(rng, ldf, init_strategy) + @test x isa AbstractVector{<:Real} + @test length(x) == LogDensityProblems.dimension(ldf) + end + end + + @testset "warning for difficult init params" begin + attempt = 0 + @model function demo_warn_initial_params() + x ~ Normal() + if (attempt += 1) < 30 + @addlogprob! -Inf + end + end + ldf = DynamicPPL.LogDensityFunction( + demo_warn_initial_params(), + DynamicPPL.getlogjoint_internal, + DynamicPPL.LinkAll(), + ) + @test_logs (:warn, r"consider providing a different initialisation strategy") Turing.Inference.find_initial_params_ldf( + Xoshiro(468), ldf, InitFromUniform() + ) + end + + @testset "errors after max_attempts" begin + @model function impossible_model() + x ~ Normal() + @addlogprob! -Inf + end + model = impossible_model() + ldf = DynamicPPL.LogDensityFunction( + model, DynamicPPL.getlogjoint_internal, DynamicPPL.LinkAll() + ) + @test_throws ErrorException Turing.Inference.find_initial_params_ldf( + Xoshiro(468), ldf, InitFromUniform() + ) + end +end + @testset "Initial parameters" begin # Dummy algorithm that just returns initial value and does not perform any sampling abstract type OnlyInit <: AbstractMCMC.AbstractSampler end diff --git a/test/mcmc/hmc.jl b/test/mcmc/hmc.jl index 5bba98bc04..424cd09b08 100644 --- a/test/mcmc/hmc.jl +++ b/test/mcmc/hmc.jl @@ -2,6 +2,7 @@ module HMCTests using ..Models: gdemo_default using ..NumericalTests: check_gdemo, check_numerical +import AbstractMCMC using Bijectors: Bijectors using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample using DynamicPPL: DynamicPPL @@ -12,7 +13,7 @@ using LinearAlgebra: I, dot, vec import Random using StableRNGs: StableRNG using StatsFuns: logistic -using Test: @test, @test_logs, @testset, @test_throws +using Test: @test, @testset, @test_throws using Turing @testset verbose = true "Testing hmc.jl" begin @@ -196,31 +197,6 @@ using Turing end end - @testset "warning for difficult init params" begin - attempt = 0 - @model function demo_warn_initial_params() - x ~ Normal() - if (attempt += 1) < 30 - @addlogprob! -Inf - end - end - - # verbose=false to suppress the initial step size notification, which messes with - # the test - @test_logs (:warn, r"consider providing a different initialisation strategy") sample( - demo_warn_initial_params(), NUTS(), 5; verbose=false - ) - end - - @testset "error for impossible model" begin - @model function demo_impossible() - x ~ Normal() - @addlogprob! -Inf - end - - @test_throws ErrorException sample(demo_impossible(), NUTS(), 5) - end - @testset "NUTS initial parameters" begin @model function f() x ~ Normal() @@ -297,12 +273,8 @@ using Turing spls = [HMC(0.1, 10), HMCDA(0.8, 0.75), NUTS(0.5), NUTS(0, 0.5)] @testset "$(spl)" for spl in spls # Construct a HMC state by taking a single step - hmc_state = Turing.Inference.initialstep( - Random.default_rng(), - gdemo_default, - spl, - DynamicPPL.VarInfo(gdemo_default); - initial_params=InitFromUniform(), + hmc_state = AbstractMCMC.step( + Random.default_rng(), gdemo_default, spl; initial_params=InitFromUniform() )[2] # Check that we can obtain the current step size @test Turing.Inference.getstepsize(spl, hmc_state) isa Float64