Skip to content

Commit

Permalink
Introduce Pad Op before tf.Split and Slice op after tf.concat ops to …
Browse files Browse the repository at this point in the history 
…handle not-divisible sharding for SPMD.

This is an alternate approach to using XLA ND Split/Concat ops. tf.Split and tf.Concat ops operate on a single dimension at a time. So the padding and slice ops are introduced accordingly.

PiperOrigin-RevId: 677052836
  • Loading branch information
@ishark @tensorflower-gardener
ishark authored and tensorflower-gardener committed Sep 21, 2024
1 parent d039ede commit 3984cff
Show file tree
Hide file tree
Showing 3 changed files with 344 additions and 58 deletions.
Loading

0 comments on commit 3984cff

Please sign in to comment.