https://github.com/adam-smnk approved this pull request. Fix itself looks fine 👍 AFAIK, the issue here was pretty generic as in type mismatch occurred during partial distribution. Won't this problem occur for any other ops? https://github.com/llvm/llvm-project/pull/150545