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.