Terrible performance of NUTS with refactored ODE API

[michaelosthege]

Description of your problem

My last-minute PR #3634 refactored the ODE API towards more explicit shapes. Unfortunately (and not detected by my benchmark) it also slowed down NUTS by about 50x compared to the original implementation by Demetri.

For reference, the approximate seconds/it/chain:

	original	refactor
NUTS	0.1	5
DEMetropolis	0.05	0.05

Also see the discussion over at #3634.

Versions and main components

• before vs. after the refactor-PR

[junpenglao]

Ouch

[michaelosthege]

When printing the parameters of every forward simulation, I made three interesting observations:

1. in the refactored implementation, NUTS has a larger tree size & depth --> more forward simulations
2. it looks like no parameter set is simulated twice --> sensitivity caching works fine(?)
3. the parameters are printed with 8 digits before & 16 digits after the refactor

So my hypothesis would be that by consistently using floatX, NUTS does more accurate leapfrogs, leading to more gradient evaluations and higher runtime.