Thanks Ricardo.
Although slightly unfortunate as the solution will probably be less generic and harder to understand 
I skimmed the NUTS/BaseHMC/ArrayStepShared code, it’s quite complex and there is relatively little documentation/examples on custom steps in pymc5.
Do you have a sense what the minimal change could possibly be?
Where should I intervene? on step? astep? _hamiltonian_step?
How will the custom step be incorporated in the pm.Model()?
Many thanks