@_sean_silva - I’d say tracing is generally a preferred way for export-y AOT-y workflows as it erases a lot of pythonism that is hard to remove with static analysis (that is a common experience with TorchScript). There’s some early work on making tracing more predictable and sound: allowing to be selective on dynamic things (e.g. varying sizes with fixed rank) and erroring out if some dynamic properties escape into python in a non-traceable way. There might be a sweet spot of mostly tracing and selective control flow at higher levels of the model code. @suo is working on some of it and can comment more
I’d be very curious in your AOT non-traceable use cases too