@DiTian @fPecc any thoughts on how to programmatically introduce annotate_sharding calls on an existing module with calls to private functions?
annotate_sharding