Avoiding automatic type promotion from F32 to F64 in NUTS

[SamuelBrand1]

I have a project for extending Distributions.jl to a new system of distribution types which use float32 float types throughout parameters and return type.

In some experiments on running toy inference problems (a log-Random walk with Poisson link) I realised that with a small number of changes DynamicHMC can be fully parametric on its float type e.g. the step size etc.

The overall goal is to have the ability to try out lifting some/all of the inference problem to compute on a GPU device. From numpyro we know that this can be advantageous in some situations and IMV having a parametric float system, as opposed to enbedded f64 assumptions in some places, is a good first step on this.

EDIT: The project is here https://github.com/SamuelBrand1/DistributionsF32 Pretty nascent still!

[tpapp]

I am open to this, PRs welcome, happy to answer questions if anything crops up.

[SamuelBrand1]

Thanks @tpapp .

I've had a go in #199 : obviously changing the API is unacceptable! I've mainly just converted some of the key structs to parametric on the float type of various scalars (e.g. stepsize). The onus should be on a user who wants a different float type to give starting position in the correct float eltype (i.e. Float32) other wise I think everything should just stay the same.

[penelopeysm]

@SamuelBrand1, I'm very much in support of adding f32 support for DynamicHMC (see also: TuringLang/Turing.jl#2212 (comment)), but I'm curious as to why you need to rewrite Distributions? It's not perfect, but you can already get a lot of f32 things out by just passing f32 things in. Sure, you need to remember to pass in f32s, but it's similar to remembering that you have to write Normal32 rather than Normal.

julia> using Distributions

julia> dist = Normal(0.0f0, 1.0f0)
Normal{Float32}(μ=0.0f0, σ=1.0f0)

julia> rand(dist)
-0.65983963f0

julia> logpdf(dist, 0.0f0)
-0.9189385f0

[SamuelBrand1]

@SamuelBrand1, I'm very much in support of adding f32 support for DynamicHMC (see also: TuringLang/Turing.jl#2212 (comment)), but I'm curious as to why you need to rewrite Distributions? It's not perfect, but you can already get a lot of f32 things out by just passing f32 things in. Sure, you need to remember to pass in f32s, but it's similar to remembering that you have to write Normal32 rather than Normal.

julia> using Distributions

julia> dist = Normal(0.0f0, 1.0f0)
Normal{Float32}(μ=0.0f0, σ=1.0f0)

julia> rand(dist)
-0.65983963f0

julia> logpdf(dist, 0.0f0)
-0.9189385f0

It depends on the underlying functions, Normal has a hard coded f64 constant in the full pdf so I suspected it of having a hidden lift to f64 in its compute. This might be wrong though.

Doing poisson needed a decent rewrite because lots of places where Int was hard-coded in, which turns out to be much more problematic than you'd expect.