Question about hierarchical model support in dynamax

[albertpod]

Hello!

I couldn't locate the discussions section, so I'm submitting this as an issue. However, I don't actually have a problem—just a question.

First, I want to commend the fantastic work on the Dynamax package. The examples are exceptionally well-crafted and very insightful—truly impressive!

I am currently benchmarking some of your examples against RxInfer.jl and am particularly impressed with Dynamax's scalability. I have a specific question regarding the support for hierarchical models in Dynamax. Is it possible to enhance the following example with additional layers of complexity? For instance, could I modify the MLP model in the nonlinear Gaussian SSM example to include higher levels of nonlinearities, such as:

• State transition:
  $$a_t = f(a_{t-1}) + w_t, \quad w_t \sim \mathcal{N}(0, 0.01 I)$$

• Parameter update:
  $$\theta_t = a_t \theta_{t-1} + q_t, \quad q_t \sim \mathcal{N}(0, 0.01 I)$$

• Observation model:
  $$y_t = h(\theta_t, x_t) + r_t, \quad r_t \sim \mathcal{N}(0, \sigma^2)$$

This is a simplified representation of the model I am considering. Thank you for your dedication and excellent work!

Cheers,
Albert