search

tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
399 stars 49 forks source link

ttnn.split should support incomplete splits #12859

Open jaykru-tt opened 14 hours ago

jaykru-tt commented 14 hours ago

ttnn.split (analogous to torch.chunk) currently doesn't have support for incomplete splits (splits for which the number of elements in the dimension on which we split is not divisible by num_splits). For RM tensors, adding support for this may be as simple as ensuring that slice handles slices similarly to torch.slice and removing the TT_FATAL at the beginning of the split op checking for divisibility of shape[dim] by num_splits.

jaykru-tt commented 14 hours ago

While we're at it, we should rename ttnn.split to ttnn.chunk for consistency with torch.