You are getting the logp in graph error because mu and sigma are random variables (because they depend on aim), but you are trying to treat them like deterministic computations. You’re not actually forced to pass observations into CustomDist, so I suppose you could concatenate mu and sigma together into a single output, then return that combined mu_sigma from your dist function. Then your observed variable could be pm.Normal(..., mu=mu_sigma[:, 0], sigma=mu_sigma[:, 1], observed=data).
That seems a bit needlessly complex though. Why is the sequence of aim (which is a normal parameterized by des and sig at each timestep) and giving it a sequence of observations is different from taking a sequence of des and `sig, using them to parameterize a sequence of normals, and instead giving that the data is different? My first instinct is that they should be completely equivalent, but I’m also worried about having to integrate over all the switches, and whether the “manual” approach will do that correctly. If you let the automatic logp inference handle it you don’t have to worry about this.