Problem conditioning on `vmap` Poisson random variables

Hi everyone,

I'm finding `oryx` a really clean approach to implementing a PPL. However, I'm confused about conditional sampling.

## Poisson process with exp random walk intensity

As an attempt to get into the structure of `oryx` I'm trying to sample from a probabilistic program which represents:

1. A hierarchical random walk (that is a random walk where the parameters are themselves random variables)
2. A further `Exp` transform on the random walk represents the intensity of a Poisson process. This is observed.
3. Inference done with `NUTS` from `blackjax`

### Code

**Dependencies**
```python
import jax
# jax.config.update("jax_enable_x64", True)

import oryx.core.ppl as ppl
import oryx.bijectors as bijectors
import oryx.distributions as tfd
import blackjax

import jax.numpy as jnp
import jax.random as random

from functools import partial
import matplotlib.pyplot as plt
```

**Prob Program**
Note that I've implemented the link as a `vmap` over `intensity` representing conditional independence of observations.

```python
@partial(jax.jit, static_argnames=["n"])
def hierarchical_random_walk_dist(n, init, step_scale):
rw_transformation = bijectors.Chain([bijectors.Shift(init), bijectors.Scale(step_scale), bijectors.Cumsum()])
return tfd.TransformedDistribution(tfd.MultivariateNormalDiag(jnp.zeros(n), jnp.ones(n)), rw_transformation)

def poisson_process(key, n, init_prior_loc, init_prior_scale, step_scale_prior):
key_poi, key_intensity, key_init, key_step = random.split(key, 4)
init = ppl.random_variable(tfd.Normal(init_prior_loc, init_prior_scale), name = "init")(key_init)
step_scale = ppl.random_variable(tfd.HalfNormal(step_scale_prior), name = "step_scale")(key_step)
intensity = ppl.random_variable(tfd.TransformedDistribution(
hierarchical_random_walk_dist(n, init, step_scale),
bijectors.Exp()),
name = "intensity")(key_intensity)
poi_keys = random.split(key_poi, n)
poi = jax.vmap(lambda ky, x: ppl.random_variable(tfd.Poisson(x), name = "poi")(ky))(poi_keys, intensity)
return poi
```

**Sample some data from model**

```python
sampler = ppl.joint_sample(poisson_process)
key_rn = random.PRNGKey(1234)
true_params = sampler(key_rn, n=100, init_prior_loc=0.0, init_prior_scale=0.05, step_scale_prior=0.25)

plt.plot(true_params['intensity'])
plt.scatter(range(len(true_params['poi'])), true_params['poi'], color='red')
plt.xlabel('time')
plt.ylabel('Intensity')
plt.title('Intensity Random Variable')
plt.show()
```


This looks reasonable.

**Inference**

I split out the observed data from the rest of the parameters

```python
true_data = true_params.pop('poi')
true_params
```
> {'init': Array(0.01620068, dtype=float32),
'intensity': Array([2.21524286e+00, 1.02979910e+00, 3.47778827e-01, 1.59668833e-01,
1.43431634e-01, 1.89372957e-01, 1.48013055e-01, 8.88996869e-02,
1.40973523e-01, 1.33820206e-01, 6.49942383e-02, 4.85427566e-02,
1.31719252e-02, 1.61807686e-02, 1.93237811e-02, 7.30752666e-03,
3.75548634e-03, 3.23897717e-03, 4.46824962e-03, 3.59713589e-03,
3.48433293e-03, 4.54167370e-03, 8.35305359e-03, 7.45324651e-03,
1.69865470e-02, 2.82925181e-03, 4.80814092e-03, 5.73506765e-03,
1.29247606e-02, 2.23501474e-02, 2.60949116e-02, 2.78504174e-02,
2.95239929e-02, 3.01535334e-02, 3.17600109e-02, 5.96645549e-02,
1.58876508e-01, 4.15319920e-01, 2.83102959e-01, 3.94434422e-01,
6.21528685e-01, 9.56910610e-01, 4.71480668e-01, 3.51778269e-01,
3.12051624e-01, 3.87135684e-01, 4.41913813e-01, 8.34466696e-01,
1.12293482e+00, 9.62291718e-01, 6.46639347e-01, 1.22468376e+00,
1.33461881e+00, 9.76860523e-01, 1.60133433e+00, 4.31086159e+00,
3.78359699e+00, 4.50091076e+00, 7.61642456e+00, 9.94997692e+00,
1.83034401e+01, 1.86841354e+01, 1.93865471e+01, 4.10644569e+01,
5.11959839e+01, 3.74023285e+01, 1.46664228e+01, 2.73789101e+01,
5.09101982e+01, 2.61694183e+01, 2.90790100e+01, 1.15916996e+01,
1.44228182e+01, 8.16150761e+00, 1.21826038e+01, 8.52718925e+00,
8.82525539e+00, 1.47077036e+01, 1.31940975e+01, 8.21146393e+00,
5.06118011e+00, 3.73972368e+00, 8.66150951e+00, 8.86765766e+00,
1.82184372e+01, 2.03960953e+01, 1.50705147e+01, 3.58565903e+01,
3.94253273e+01, 1.51045656e+01, 1.46200066e+01, 1.30218935e+01,
8.60846615e+00, 6.86474085e+00, 5.52572966e+00, 6.91005898e+00,
4.49717140e+00, 2.01037908e+00, 2.75382376e+00, 3.28753996e+00], dtype=float32),
'step_scale': Array(0.5200044, dtype=float32)}

Then do the usual `blackjax` approach to sampling (based on their example of using `oryx`)

```python
def logdensity_fn(params):
theta = dict(params, poi = true_data)
return ppl.joint_log_prob(poisson_process)(theta, n=100, init_prior_loc=0.0, init_prior_scale=0.05, step_scale_prior=0.25)

ll = logdensity_fn(true_params)
# Array(-50.80758, dtype=float32)

# Warmup
inference_key = jax.random.PRNGKey(12)
rng_key, warmup_key = jax.random.split(inference_key)
adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
(last_state, parameters), _ = adapt.run(warmup_key, true_params, 1000)
kernel = blackjax.nuts(logdensity_fn, **parameters).step

# Sampling

def inference_loop(rng_key, kernel, initial_state, num_samples):
def one_step(state, rng_key):
state, info = kernel(rng_key, state)
return state, (state, info)

keys = jax.random.split(rng_key, num_samples)
_, (states, infos) = jax.lax.scan(one_step, initial_state, keys)

return states, infos

rng_key, sample_key = jax.random.split(rng_key)
states, infos = inference_loop(sample_key, kernel, last_state, 2000)
```

### Issues

1. The main issue is that this silent fails to sample from the posterior (or I'm not understanding the sample structure):

```python
plt.figure(figsize=(12, 6))
plt.plot(true_params['intensity'], label='True Intensity', color='blue')
for i in range(100): # Plotting the first 10 sampled intensities for clarity
plt.plot(states.position['intensity'][i], alpha=0.5)
plt.xlabel('Time')
plt.ylabel('Intensity')
plt.title('True Intensity vs Sampled Intensities')
plt.legend()
plt.show()
```


2. Warnings about `f32` conversion e.g.
> UserWarning: Explicitly requested dtype requested in zeros is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
minval = minval + np.zeros([1] * final_rank, dtype=dtype)

Which suggests that the underlying Poisson distribution is struggling but...

3. Errors if enabling `f64` conversion, e.g. if the conversion to double precision is allow then the model fails at `joint_sample` with error:
> TypeError: Tensor conversion requested dtype for array with dtype float64: Tracedwith

### Steps forward

I don't have a huge amount of JAX/oryx experience, therefore, it would be great if someone could point out if I've made a glaring error. Or if there is some kind of issue with `joint_log_prob` in combination with `Poisson` or the way I've implemented the poisson link.

Reproduce the provided `poisson_process` example and inspect the `ppl.joint_sample` and `ppl.joint_log_prob` entry points, especially the `vmap` over Poisson variables. Compare the posterior state structure with the observed `poi` data and investigate the reported float32/float64 failures. Done means the conditioning behavior and dtype path are understood and the reported sampling problem has a verified outcome.

Written by the indexing model from the issue text.