lightning-thunder icon indicating copy to clipboard operation
lightning-thunder copied to clipboard

[DTensor] Refactor to re-use prims

Open kshitij12345 opened this issue 1 month ago • 1 comments

Investigate: Try to decrease duplication of prims - reuse the existing prims one by one to support propagating DTensor layout and providing the executors information about the supported types (Tensor, DTensor) which the executor opt in or out of. (See comment)

kshitij12345 avatar Nov 10 '25 08:11 kshitij12345

Note that you're converting a 2xN problem into a Nx2 problem by pushing dtensor support into prims. In the olden days, we decomposed ops to factor out common patterns, maybe this could still be a route.

t-vi avatar Nov 11 '25 11:11 t-vi