| 141 | } |
| 142 | |
| 143 | XlaOp TorchGather(XlaOp input, XlaOp index, int64 dim, bool sparse) { |
| 144 | XlaBuilder* builder = input.builder(); |
| 145 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 146 | TF_ASSIGN_OR_RETURN(Shape index_shape, builder->GetShape(index)); |
| 147 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 148 | if (ShapeUtil::ElementHasBitWidth(index_shape, 64) && |
| 149 | input_shape.dimensions(dim) < std::numeric_limits<uint32>::max()) { |
| 150 | index = ConvertElementType(index, U32); |
| 151 | index_shape.set_element_type(U32); |
| 152 | } |
| 153 | if (index_shape.rank() == 1) { |
| 154 | return TorchIndexSelect(input, index, 0); |
| 155 | } |
| 156 | if (!sparse) { |
| 157 | std::vector<int64> index_broadcast_dims; |
| 158 | std::vector<int64> input_broadcast_dims; |
| 159 | std::vector<int64> sizes; |
| 160 | for (int64 i = 0; i < index_shape.rank(); ++i) { |
| 161 | if (i < dim) { |
| 162 | input_broadcast_dims.push_back(i); |
| 163 | index_broadcast_dims.push_back(i); |
| 164 | } else if (i == dim) { |
| 165 | sizes.push_back(input_shape.dimensions(i)); |
| 166 | input_broadcast_dims.push_back(i); |
| 167 | index_broadcast_dims.push_back(i + 1); |
| 168 | } else { |
| 169 | input_broadcast_dims.push_back(i + 1); |
| 170 | index_broadcast_dims.push_back(i + 1); |
| 171 | } |
| 172 | sizes.push_back(index_shape.dimensions(i)); |
| 173 | } |
| 174 | auto mask = Eq( |
| 175 | BroadcastInDim(index, sizes, index_broadcast_dims), |
| 176 | Iota(builder, ShapeUtil::MakeShape(index_shape.element_type(), sizes), |
| 177 | dim)); |
| 178 | auto masked_input = Select( |
| 179 | mask, BroadcastInDim(input, sizes, input_broadcast_dims), |
| 180 | Zeros(builder, |
| 181 | ShapeUtil::MakeShape(input_shape.element_type(), sizes))); |
| 182 | return Reduce(masked_input, Zero(builder, input_shape.element_type()), |
| 183 | CreateScalarIdentityWithZeroComputation( |
| 184 | input_shape.element_type(), builder), |
| 185 | {dim}); |
| 186 | } |
| 187 | |
| 188 | ShapeUtil::AppendMajorDimension(1, &index_shape); |
| 189 | std::vector<XlaOp> to_concat; |
| 190 | |
| 191 | to_concat.reserve(input_shape.rank()); |
| 192 | for (int64 i = 0; i < input_shape.rank(); ++i) { |
| 193 | if (i == dim) { |
| 194 | to_concat.push_back(Reshape(index, index_shape.dimensions())); |
| 195 | } else { |
| 196 | to_concat.push_back(Iota(builder, index_shape, i)); |
| 197 | } |
| 198 | } |
| 199 | XlaOp gather_indices = ConcatInDim(builder, to_concat, input_shape.rank()); |
| 200 | std::vector<int64> slice_sizes(input_shape.rank(), 1); |