fix: diffusion sampling

[ssmmnn11]

• fix diffusion sampling after multi-dataset integration and multi out
• sampler operations consistently in dtype
• improve tests

[JPXKQX]

Thanks for fixing this Simon! I think that at some point we need to update the diffusion section of our docs to show/explain how to choose these different diffusion samplers but I'm happy to do this in a later PR. Feel free to merge when the integration tests pass.