[AFQ] Optimize tensor_flatten for runtime

[IvanKobzarev]

Stack from ghstack (oldest at bottom):

• -> [AFQ] Optimize tensor_flatten for runtime #1951

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1951

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Unrelated Failure

As of commit b72c5fa with merge base 923242e ():

NEW FAILURE - The following job has failed:

• .github/workflows/float8nocompile_test.yaml (gh)

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh) (trunk failure)
  ModuleNotFoundError: No module named 'importlib_metadata'

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

please change the branch to main before merging