Allow For An Interface Without a NumPyro Dependence

[DanWaxman]

There may be desires in some settings to avoid a direct NumPyro dependency, dealing with the tracing, etc.

[DanWaxman]

Thankfully, this should be rather doable due to some recent refactoring that was required for hierarchical sampling, anyways.

Filtering

For example, all filtering integrations now come with a pure JAX implementation; for example, in cuthbert integrations:

dynestyx/dynestyx/inference/integrations/cuthbert/discrete.py

Lines 163 to 177 in c70d726

	def compute_cuthbert_filter(
	dynamics: DynamicalModel,
	filter_config: BaseFilterConfig,
	key: jax.Array | None = None,
	*,
	obs_times: jax.Array,
	obs_values: jax.Array,
	ctrl_times=None,
	ctrl_values=None,
	):
	"""Pure-JAX cuthbert filter computation (no numpyro side-effects).
	Returns:
	tuple: (marginal_loglik, states) where states are aligned to obs_times.
	"""

The corresponding NumPyro code that surrounds that simply runs the compute_cuthbert_filter path, then records some numpyro.deterministic sites + numpyro.factor statements:

dynestyx/dynestyx/inference/integrations/cuthbert/discrete.py

Lines 238 to 287 in c70d726

	def run_discrete_filter(
	name: str,
	dynamics: DynamicalModel,
	filter_config: BaseFilterConfig,
	key: jax.Array | None = None,
	*,
	obs_times: jax.Array,
	obs_values: jax.Array,
	ctrl_times=None,
	ctrl_values=None,
	**kwargs,
	) -> list[dist.Distribution]:
	"""Run discrete-time filter via cuthbert (Kalman, Taylor KF, EnKF, PF).
	Returns:
	list[dist.Distribution]: Filtered state distributions at each obs time.
	"""
	obs_len = int(obs_values.shape[0])
	if obs_len == 0:
	return []
	marginal_loglik, states = compute_cuthbert_filter(
	dynamics,
	filter_config,
	key,
	obs_times=obs_times,
	obs_values=obs_values,
	ctrl_times=ctrl_times,
	ctrl_values=ctrl_values,
	)
	record_kwargs = _config_to_record_kwargs(filter_config)
	numpyro.factor(f"{name}_marginal_log_likelihood", marginal_loglik)
	numpyro.deterministic(f"{name}_marginal_loglik", marginal_loglik)
	if isinstance(filter_config, PFConfig):
	_add_sites_pf(name, states, record_kwargs)
	particles = states.particles
	if particles.ndim == 2:
	particles = particles[..., None]
	return particles_to_delta_mixtures(particles, states.log_weights)
	else:
	_add_sites_gaussian_filter(name, states, record_kwargs)
	mean = states.mean
	chol_cov = states.chol_cov
	cov = covariance_from_cholesky(chol_cov)
	return [
	dist.MultivariateNormal(mean[i], covariance_matrix=cov[i])
	for i in range(mean.shape[0])
	]

Similar results hold for CD-Dynamax integrations.

Simulators

Simulators are a bit more messy, but I think still doable.

SDE simulator would be pretty straightforward. It mostly comprises a pure-JAX solve_sde call + a rather simple numpyro-based emit_observations call:

dynestyx/dynestyx/simulators.py

Lines 819 to 851 in c70d726

	def _sim_one_trajectory(key: Array, x0: Array) -> tuple[Array, Array]:
	"""Simulate one SDE trajectory and its emissions."""
	k_solve, k_obs = jr.split(key, 2)
	states_sol = solve_sde(
	source=self.source,
	dynamics=dynamics,
	t0=t0,
	saveat_times=times,
	x0=x0,
	control_path_eval=control_path_eval,
	diffeqsolve_settings=self.diffeqsolve_settings,
	key=k_solve,
	tol_vbt=self.tol_vbt,
	)
	emissions = _emit_observations(
	name, dynamics, states_sol, times, None, control_path_eval, key=k_obs
	)
	return states_sol, emissions
	prng_key = numpyro.prng_key()
	if prng_key is None:
	raise ValueError("PRNG key required for simulation")
	with numpyro.plate(f"{name}_n_simulations", n_sim):
	initial_state = numpyro.sample(f"{name}_x_0", dynamics.initial_condition)
	keys = jr.split(prng_key, n_sim)
	states, emissions = jax.vmap(_sim_one_trajectory)(
	keys, jnp.asarray(initial_state)
	)
	return {
	"times": _tile_times(times, n_sim),
	"states": states,
	"observations": emissions,
	}

This _emit_observations is really not doing anything special NumPyro-wise:

dynestyx/dynestyx/simulators.py

Lines 591 to 630 in c70d726

	def _emit_observations(
	name: str,
	dynamics,
	states: Array,
	times: Array,
	obs_values: Array | None,
	control_path_eval: Callable[[Array], Array | None],
	key=None,
	) -> Array:
	"""Emit observations via numpyro.sample (conditioning) or dist.sample (vmap)."""
	ctrl = control_path_eval if control_path_eval is not None else (lambda t: None)
	T = len(times)
	if key is not None:
	obs_keys = jr.split(key, T)
	def _obs_step(t_idx):
	x_t = states[t_idx]
	t = times[t_idx]
	u_t = ctrl(t)
	obs_dist = dynamics.observation_model(x=x_t, u=u_t, t=t)
	return obs_dist.sample(obs_keys[t_idx])
	return jax.vmap(_obs_step)(jnp.arange(T))
	else:
	def _step(carry, t_idx):
	x_t = states[t_idx]
	t = times[t_idx]
	u_t = ctrl(t)
	obs_t = _get_val_or_None(obs_values, t_idx)
	y_t = numpyro.sample(
	f"{name}_y_{t_idx}",
	dynamics.observation_model(x=x_t, u=u_t, t=t),
	obs=obs_t,
	)
	return carry, y_t
	_, observations = nscan(_step, None, jnp.arange(T))
	return observations

ODE simulation is a similar story.

The more painful story will be discrete-time simulation. To be honest, this is kind of a mess right now, from trying to support multiple different features (perfect observations, hierarchical simulation, batched simulation, and normal simulation). It contains a pure-JAX scan:

dynestyx/dynestyx/simulators.py

Lines 1012 to 1047 in c70d726

	# n_simulations > 1: vmapped pure-JAX loop; avoid numpyro.sample in vmap body.
	# can't do obs= conditioning with the lax.scan.
	if n_sim > 1:
	with numpyro.plate(f"{name}_n_simulations", n_sim):
	initial_state = numpyro.sample(
	f"{name}_x_0", dynamics.initial_condition
	)
	prng_key = numpyro.prng_key()
	if prng_key is None:
	raise ValueError("PRNG key required for n_simulations > 1")
	keys = jr.split(prng_key, n_sim)
	def _sim_one_trajectory(key, x0):
	"""Simulate one discrete trajectory and its emissions."""
	keys_t = jr.split(key, T)
	def _step(carry, t_idx):
	x_prev = carry
	k_trans, k_obs = jr.split(keys_t[t_idx], 2)
	t_next, u_next, trans_dist = _step_dists(x_prev, t_idx)
	x_t = trans_dist.sample(k_trans)
	y_t = dynamics.observation_model(x=x_t, u=u_next, t=t_next).sample(
	k_obs
	)
	return x_t, (x_t, y_t)
	u_0 = _get_val_or_None(ctrl_values, 0)
	y_0 = dynamics.observation_model(x=x0, u=u_0, t=times[0]).sample(
	keys_t[0]
	)
	_, (scan_states, scan_obs) = jax.lax.scan(_step, x0, jnp.arange(T - 1))
	states = jnp.concatenate([jnp.expand_dims(x0, 0), scan_states], axis=0)
	observations = jnp.concatenate(
	[jnp.expand_dims(y_0, 0), scan_obs], axis=0
	)
	return states, observations

But it also contains a lot of other stuff... I think this will be a major pain point that we have to deal with.

Other Caveats

There are lots of places where we use the numpyro PRNG key machinery, since it's pretty useful. That said, there's no reason we can't roll out our own machinery to replace it. In particular, I would propose something like dynestyx.prng_key() that falls back to numpyro.prng_key() if PYRO_STACK exists and has a seed handler in it, and otherwise, generate a global random seed objax-style.

General Design

How to use this numpyro-free pathway is an open question, to me. One option that makes sense is to consider a parallel pathway to dynestyx.sample(...), e.g., dynestyx.infer(...) instead.

Another question is how to avoid lots of duplicate code between NumPyro and non-NumPyro code. One option is, where we have pathways that only use numpyro.sample(...), we roll out a background-agnostic dynestyx._sample_tensor(...) or something, that checks if a numpyro inference context is active, uses numpyro.sample if so, and uses hand-rolled prob.sample statements otherwise. I guess this is a bit of PyMC route, which has dangers (that codebase is very, very complex), so we would need to be very careful about scope creep of this backend-agnostic interface.

[theorashid]

Aiming to attempt this in 2 weeks when back from holidays. Starting with a draft PR for Filter, ignore Simulator for now

[theorashid]

Sorry, work has been hectic since I've been back. I'll get to this next week.

[theorashid]

I've started on this today. I wrote a version that is

filt = Filter(filter_config=KFConfig(...), register_numpyro=False)
with filt:
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)
loss = -filt.marginal_loglik # use in optax or whatever

# with numpyro:
with Filter(filter_config=KFConfig(...)):
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)

But I think this is still too coupled because Filter still imports and knows about numpyro_sites and it does not need to. I'm going to try a version more like what I had written before which is

filt = Filter(filter_config=KFConfig(...))
with filt:
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)
loss = -filt.marginal_loglik

# with numpyro:
with filt, NumpyroTrace(filt):
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)

Which breaks a lot more things but I think it's better in terms of composability and separation. So I'll continue this next week (maybe over the long weekend because I'm just in Seattle and this is fun).

[DanWaxman]

Hey Theo, thanks so much!!

I agree that the first option doesn't break as cleanly as one would like. In particular, I think most functions now break pretty cleanly into "places where work happens" and "post-hoc NumPyro registration" -- this was necessary to implement plating, since we want to vmap over the "places where work happens" parts but not over anything NumPyro.

The second is along what I was imagining, but I also don't want to have another obligatory handler for whenever one wants to use NumPyro (which seems like a fairly common use case for us, though not the only one). Do you (or maybe @mattlevine22) have thoughts on introducing something like dsx.infer(...) or dsx.conditionfor the "NumPyro-free" case, and keep dsx.sample(...) as the NumPyro case? It seemed like "sample" as a function name was confusing to people anyways.

Note that this wouldn't change your implementation very much -- dsx.sample could just be an alias for with NumpyroTrace(): dsx.condition(...). In this case, I think a nice design pattern would be: infer returns a dictionary returns of information, and NumpyroTrace is a simple handler that just runs log_numpyro_values(fwd(...)). This also simplifies your proposal a bit, as you no longer have to wrap the filter implementation in NumpyroTrace(...), i.e.,

# proposal from before
with filt, NumpyroTrace(filt):
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)

# after, no NumPyro at all
with filt:
    dsx.infer("f", dynamics, obs_times=t, obs_values=y) 
    # or maybe dsx.condition(...)?

# after
with filt, NumpyroTrace(): # <--- no need to pass in filt
    dsx.infer("f", dynamics, obs_times=t, obs_values=y)
# or equivalently 
with filt:
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)

This is a bit less general (I guess in your proposal, you could attach a NumPyro trace to a filter but not a simulator), though I think the all-or-nothing approach is alright.

Let me know what you think! Happy to discuss more -- I'm plenty sure there's problems with what I propose, but I do think the PPL interface really benefits from as much simplicity as possible.

[mattlevine22]

Yeah I agree with Dan that the 2nd option is better, but would prefer the default behavior of with Filter... dsx.sample(...) does not change from user-perspective. I also agree that the non-Numpyro usage should probably NOT use dsx.sample(), as it gets overloaded this way. Instead, dsx.infer or dsx.condition seem reasonable to me.

My vision is that dsx.sample actually just runs something like foo = dsx.condition(...); dsx.add_sites(foo) (and that might mean that we have an interpretation for condition and an interpretation for add_sites, rather than the current direct interpretation of sample).

Relatedly, in the non-numpyro case, I think it would be best to simply return an object, rather than attach things to filt.

This is also exciting because one can imagine using this to plug into other PPLs (assuming they support JAX), e.g. PyMC

[theorashid]

Alright, I've got something like this but I'm going to spend some more time on it and launch a draft PR.

def sample(...):
    result = infer(...)
    if isinstance(result, InferResult) and result._register_numpyro_sites is not None:
        result._register_numpyro_sites(name)
    return result

def test_infer_does_not_register_numpyro_sites():
    from numpyro.handlers import seed, trace

    obs_times, obs_values = _make_data()
    dynamics = _make_lti_dynamics(0.5)

    with trace() as tr, seed(rng_seed=jr.PRNGKey(0)):
        with Filter(filter_config=KFConfig(filter_source="cuthbert")):
            result = dsx.infer("f", dynamics, obs_times=obs_times, obs_values=obs_values)

    assert isinstance(result, InferResult)
    assert result.marginal_loglik is not None
    assert "f_marginal_loglik" not in tr
    assert "f_marginal_log_likelihood" not in tr

def test_sample_registers_numpyro_sites():
    from numpyro.handlers import seed, trace

    obs_times, obs_values = _make_data()
    dynamics = _make_lti_dynamics(0.5)

    with trace() as tr, seed(rng_seed=jr.PRNGKey(0)):
        with Filter(filter_config=KFConfig(filter_source="cuthbert")):
            dsx.sample("f", dynamics, obs_times=obs_times, obs_values=obs_values)

    assert "f_marginal_loglik" in tr
    assert jnp.isfinite(tr["f_marginal_loglik"]["value"])

Smoother I could probably add. But getting the numpyro.deterministic out of Simulator might break a few more things.

So I'll just start with Filter draft PR, then let's discuss the design and expand once we've settled on it.

[mattlevine22]

This looks great to me!

One question---why is dsx.infer inside a numpyro trace in the test? I would not expect this to be a typical use case.

[theorashid]

Yeah I don't expect anyone to do dsx.infer inside trace in practice. The test was to make sure the numpyro artifacts produced by _register_numpyro_sites were not in the trace for infer but are there for sample.