We should add some tests checking that gradients computed by our rrules are correct. For this, we can use [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl) or [ChainRulesTestUtils.jl](https://github.com/JuliaDiff/ChainRulesTestUtils.jl).