[TensorExpr] Reland: "Simplify index expressions constructed in loop flattening. Fixes #51173" (#53861)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53861
Replaced the iterators in the for-loops with integer index variables due to
overflow when handling empty vectors.
Test Plan: Imported from OSS
Reviewed By: navahgar
Differential Revision: D26998894
Pulled By: huiguoo
fbshipit-source-id: a1f6475c8ba123968ef7247b4f6f38edbf24b9ef
