Add type-stability checks via `DispatchDoctor.jl`

[avik-pal]

1. We can't apply it to the entire Lux module. Some of the code exists for debugging and is intentionally type unstable.
2. I like the default of erroring, but that will be too breaking. Maybe we can consider that for v0.6. We will warn if the code is type unstable for now.

Repository: https://github.com/MilesCranmer/DispatchDoctor.jl