@RGon Yeah, I think that’s the end goal. It’s possible to use MTK that way right now (I’ve done it to some approximation) but it requires a fair amount of hacking. For starters, you need to avoid using symbolic indexing in the loss function, which means that you need to have a loss function factory that finds the index of the state variables that you want to derive the loss from, and then capture that indexer in the loss function. (See the FAQ). Also, if you need an observed variable, you have to get the observation function and capture it too. Finally, if you want some parameters to be adjustable & others not, you need to do more indexing tricks to interpolate the values of the adjustable subset into the full parameter vector & fill the rest with the default values. More indexing is required if you want to override the default value of a parameter.
I think the plan for the future involves fixing at least some of those inconveniences. For instance, to toggle the adjustability of a parameter, there’s going to be a metadata flag that will tell optimizers not to adjust a parameter. There may also be some work on getting observed variables more efficiently. IDK if that will help with automatic differentiation though.