mxtensor: add index select support

[vkuzo]

Summary:

this is needed for stitching 2d expert weights to 3d in vllm

Test Plan:

pytest test/prototype/mx_formats/test_mx_tensor.py -s -k test_index_select

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• enroll mxtensor in vllm integration tests #3081
• mxtensor: switch to AOBaseTensor dispatch #3080
• -> mxtensor: add index select support #3079

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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