| 246 | } |
| 247 | |
| 248 | XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim, int64 batch_dims) { |
| 249 | XlaBuilder* builder = input.builder(); |
| 250 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 251 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 252 | TF_ASSIGN_OR_RETURN(Shape index_shape, builder->GetShape(index)); |
| 253 | if (dim < batch_dims) { |
| 254 | return InvalidArgument( |
| 255 | "Gather dim must be greater than or equal to the number of batch " |
| 256 | "dims"); |
| 257 | } |
| 258 | if (ShapeUtil::ElementHasBitWidth(index_shape, 64) && |
| 259 | input_shape.dimensions(dim) < std::numeric_limits<uint32>::max()) { |
| 260 | index = ConvertElementType(index, U32); |
| 261 | index_shape.set_element_type(U32); |
| 262 | } |
| 263 | std::vector<int64> slice_sizes = SpanToVector(input_shape.dimensions()); |
| 264 | GatherDimensionNumbers gather_dnums; |
| 265 | gather_dnums.set_index_vector_dim(index_shape.rank()); |
| 266 | if (batch_dims > 0) { |
| 267 | ShapeUtil::AppendMajorDimension(1, &index_shape); |
| 268 | std::vector<XlaOp> to_concat; |
| 269 | to_concat.reserve(batch_dims + 1); |
| 270 | for (int64 batch_dim = 0; batch_dim < batch_dims; ++batch_dim) { |
| 271 | to_concat.push_back(Iota(builder, index_shape, batch_dim)); |
| 272 | } |
| 273 | to_concat.push_back(Reshape(index, index_shape.dimensions())); |
| 274 | index = ConcatInDim(builder, to_concat, gather_dnums.index_vector_dim()); |
| 275 | } |
| 276 | for (int64 i = 0; i < input_shape.rank(); ++i) { |
| 277 | if (i < batch_dims || i == dim) { |
| 278 | slice_sizes[i] = std::min<int64>(slice_sizes[i], 1); |
| 279 | gather_dnums.add_collapsed_slice_dims(i); |
| 280 | gather_dnums.add_start_index_map(i); |
| 281 | } else { |
| 282 | if (i < dim) { |
| 283 | gather_dnums.add_offset_dims(i); |
| 284 | } else { |
| 285 | gather_dnums.add_offset_dims(i + gather_dnums.index_vector_dim() - |
| 286 | (1 + batch_dims)); |
| 287 | } |
| 288 | } |
| 289 | } |
| 290 | return Gather(input, index, gather_dnums, slice_sizes); |
| 291 | }); |
| 292 | } |
| 293 | |
| 294 | } // namespace xla |