Support `Float32` and other precisions

[yebai]

The assumption of only supporting Float64 is too strong. Are there any good reasons?

I am labelling this as bugs to signal its priority.

[willtebbutt]

The reason is a lack of testing infrastructure which can properly handle Float32 -- with multiple precisions involved, it becomes a little bit more tricky to properly determine appropriate accuracy thresholds for finite differencing etc.

I've removed the bug label, in favour of enhancement + high priority, because bug isn't an accurate description of what's going on here.

[willtebbutt]

@yebai and @Red-Portal discovered that we're going to need a rule to Core.Intrinsics.fptrunc is going to need a conversion rule + a rule -- see TuringLang/AdvancedVI.jl#71

[willtebbutt]

@yebai plan for this below:

We plan to support testing functions involving uniform-precision data. For example,

f(x::Float64, y::Float64)
f(x::Float32, y::Float32)

should be fine, but we do not intend to add support for functions containing a range of precisions. For example,

f(x::Float32, y::Float64)

would not be testable.

In order to support this, we would need to:

1. produce some functionality which checks that only a single precision has been used in the arguments to a test case,
2. add functionality to choose the finite differencing step size and the various accuracy thresholds, based on whether Float16 / Float32 / Float64 is used.

Once this is in place, we will need to

1. add tests for Float32 all over the place and fix any problems which appear, and
2. add support for specific functions, such as Core.Intrinsics.fptrunc, which we do not currently support.

[yebai]

That sounds good to me; maybe we could promote low-precision types to the highest ones in arguments as a default?

[willtebbutt]

I think I'd rather just provide an informative error message, in order to avoid users thinking they're testing one thing, but actually testing another -- I can easily imagine a situation arising in which a user accidentally has a Float64 somewhere in their input, and then the whole thing gets silently promoted to Float64, and then they're testing something other than what they meant to test.

[willtebbutt]

@penelopeysm discovered in TuringLang/docs#559 (comment) that Mooncake's current rule for matrix-matrix multiplication in LuxLib doesn't successfully handle the case that the two input arrays contain numbers at different precisions.

It really shouldn't be too hard to handle this properly -- we would just need to re-write the current from_rrule implementations of the various variants of matrix-matrix multiplication found around here to have proper Mooncake rules which

1. define is_primitive for arrays whose elements are subtypes of IEEEFloat, and
2. define the rrule!s in such a way that ensures the correct element-type is adhered to.
3. add a method of rrule!! which is a catch-all for all other element types, which always errors with some kind of sensible error message that users can make use of to know how to modify their code.

Note: this is also a great opportunity to ensure excellent performance in Lux.jl -- the current implementations of the rules involve more allocations than are really needed, because we do not increment the gradients in-place.

[willtebbutt]

I'm marking this as done, because I believe #414 basically solves it . We can re-open if there is a reason to.

[yebai]

It really shouldn't be too hard to handle this properly -- we would just need to re-write the current from_rrule implementations of the various variants of matrix-matrix multiplication found around here to have proper Mooncake rules which

1. define is_primitive for arrays whose elements are subtypes of IEEEFloat, and
2. define the rrule!s in such a way that ensures the correct element-type is adhered to.
3. add a method of rrule!! which is a catch-all for all other element types, which always errors with some kind of sensible error message that users can make use of to know how to modify their code.

Has this been addressed in previous PRs?