What would it take to support Distributions with pytree-valued samples in numpyro?
The weiner first passage time distribution produces a discrete integer choice and a floating point response time. These can be concatenated together into a single floating point array of size (…, 2) but it’s a bit odd. The timed-race diffusion model is another distribution like that. In general, it seems like you should be able to specify a log probability for any pytree, not just a single array.