Searched refs:index_broadcast_dims (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | slicing.cc | 172 std::vector<int64> index_broadcast_dims; in TorchGather() local
178 index_broadcast_dims.push_back(i); in TorchGather()
182 index_broadcast_dims.push_back(i + 1); in TorchGather()
185 index_broadcast_dims.push_back(i + 1); in TorchGather()
190 BroadcastInDim(index, sizes, index_broadcast_dims), in TorchGather()
232 std::vector<int64> index_broadcast_dims; in TorchScatterDense() local
236 index_broadcast_dims.push_back(i); in TorchScatterDense()
241 index_broadcast_dims.push_back(i + 1); in TorchScatterDense()
246 Eq(BroadcastInDim(index, sizes, index_broadcast_dims), in TorchScatterDense()
250 Select(mask, BroadcastInDim(src, sizes, index_broadcast_dims), in TorchScatterDense()
|