As discussed on Slack, it would be good to have an additional specialization of logdensity that takes NamedTuple-valued inputs. The Turing main repo already implements this as boilerplate for AdvancedMH, so it would be nice to have an official unified implementation that also supports AD.
Right now, it seems to me that adding a specialization to unflatten would do the job. (And maybe change the function's name to something else since it is not doing unflatten anymore?)