Hey!
We are looking into a generalized version of that extension system.
But DTensor has a special one that you can use via pytorch/torch/distributed/_tensor/_sharding_prop.py at 03440a1c13caf5fdb060245483c0056a5b43627f · pytorch/pytorch · GitHub
3 Likes