Add Annotated Event Dimensions

[DanWaxman]

A persistent pain point in plating has been the lack of named event dimensions, leading to ambiguities between batch dimensions and event dimensions. It would be wonderful to resolve this, so that we avoid special case pathways like

dynestyx/dynestyx/utils.py

Lines 113 to 140 in 6cba16c

	# Whitelist of built-in model fields whose trailing axis is a vector event axis.
	#
	# A shared vector whose length happens to equal a plate size is otherwise
	# ambiguous (is `(N,)` a per-member scalar or a shared length-N vector?). The
	# conservative read is "shared," and `_leaf_is_plate_batched` skips rank-1
	# suffixes by default. This whitelist opts specific built-in fields back in:
	# for these paths, a rank-1 suffix is *known* to be a vector event axis, so
	# `(N, d)` should be treated as plate-batched even when `d == 1`.
	#
	# Pinned by:
	# tests/test_hierarchical_smokes.py::test_unbatched_vector_fields_matching_plate_size_remain_shared
	#
	# To extend: add the (parent_field, ..., leaf_field) tuple here and add a
	# matching smoke test exercising both the shared and plate-batched cases.
	def _is_known_vector_field(path) -> bool:
	"""Return True for built-in leaves whose final axis is a vector event axis."""
	names = _path_field_names(path)
	# `Discretizer` wraps the original continuous-time evolution in a `cte` field
	# of `EulerMaruyamaGaussianStateEvolution`, so a drift bias that lived at
	# `state_evolution.drift.b` moves to `state_evolution.cte.drift.b`. Drop that
	# internal wrapper segment so the same whitelist matches discretized models.
	names = tuple(name for name in names if name != "cte")
	if len(names) >= 2 and names[-2:] in {
	("state_evolution", "bias"),
	("observation_model", "bias"),
	}:
	return True
	return len(names) >= 3 and names[-3:] == ("state_evolution", "drift", "b")

or

dynestyx/dynestyx/inference/plate_utils.py

Lines 16 to 45 in 70c1ac2

	def _make_plate_in_axes(tree, plate_shapes: tuple[int, ...]):
	"""Build ``in_axes`` for leaves whose leading dims match active plates.
	All numpyro distributions are treated as opaque leaves with ``in_axes=None``.
	A plate-batched ``initial_condition`` is *not* sliced by ``vmap`` here; the
	batched filter/smoother dispatch rebuilds it per member from the clean
	original via ``_slice_dist_for_plate_member``. This avoids ``vmap`` leaving a
	stale ``batch_shape`` in the distribution's static aux-data (which would make
	``.mean`` / ``.sample`` / ``.log_prob`` re-expand to the full plate shape).
	"""
	def _axis(path, leaf):
	if isinstance(leaf, numpyro.distributions.Distribution):
	return None
	# Only constant-coefficient diffusions are opaque leaves (see
	# ``_is_opaque_plate_leaf``); a callable coefficient is recursed into, so
	# its array fields are vmapped generically by the branch below.
	if isinstance(leaf, Diffusion):
	return (
	0
	if _diffusion_coefficient_is_plate_batched(leaf, plate_shapes)
	else None
	)
	return 0 if _leaf_is_plate_batched(leaf, plate_shapes, path=path) else None
	return jax.tree_util.tree_map_with_path(
	_axis,
	tree,
	is_leaf=_is_opaque_plate_leaf,
	)

This is essentially a named dimension; what we are doing is trying to detect the event dimension so that we can properly determine if the total shape is plate_shapes + event_shapes. It might be a good fit for, e.g., effectful named dimensions?

[DanWaxman]

According to @jfeser, the right place to look for in effectful will be the bind_dims/unbind_dims interface.