Searched refs:gather_dnums (Results 1 – 1 of 1) sorted by relevance
154 GatherDimensionNumbers gather_dnums; in TorchGather() local155 gather_dnums.set_index_vector_dim(input_shape.rank()); in TorchGather()157 gather_dnums.add_collapsed_slice_dims(i); in TorchGather()158 gather_dnums.add_start_index_map(i); in TorchGather()160 return Gather(input, gather_indices, gather_dnums, slice_sizes); in TorchGather()171 GatherDimensionNumbers gather_dnums; in TorchIndexSelect() local174 gather_dnums.add_offset_dims(i); in TorchIndexSelect()177 gather_dnums.set_index_vector_dim(index_shape.rank()); in TorchIndexSelect()178 gather_dnums.add_collapsed_slice_dims(dim); in TorchIndexSelect()179 gather_dnums.add_start_index_map(dim); in TorchIndexSelect()[all …]