It would be nice to support python scalars as start or end. Similar work has been done for torch.where cc @gchanan @mruberry @nairbv