mxtensor: support clone

[vkuzo]

Summary:

This is needed for saving HF models with mxfp4 weights to disk.

Test Plan:

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

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
• mxtensor: add serialization support #3078
• mxtensor: delete to_copy override #3072
• -> mxtensor: support clone #3070

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.

[danielvegamyhre]

Why not just avoid the if statement with clone_fn = lambda x: x.clone(**kwargs)? Do we need to explicitly exclude other kwargs and only include memory format if specified, for some reason?

[vkuzo]

this is copy-pasta from

ao/torchao/prototype/mx_formats/nvfp4_tensor.py

Line 346 in 7690612

def nvfp4_clone(func, types, args, kwargs):

, IMO let's separate cleanup of this into a separate PR