Searched refs:to_concat (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | slicing.cc | 204 std::vector<XlaOp> to_concat; in TorchGather() local
206 to_concat.reserve(input_shape.rank()); in TorchGather()
209 to_concat.push_back(Reshape(index, index_shape.dimensions())); in TorchGather()
211 to_concat.push_back(Iota(builder, index_shape, i)); in TorchGather()
214 XlaOp gather_indices = ConcatInDim(builder, to_concat, input_shape.rank()); in TorchGather()
283 std::vector<XlaOp> to_concat; in TorchIndexSelect() local
284 to_concat.reserve(batch_dims + 1); in TorchIndexSelect()
286 to_concat.push_back(Iota(builder, index_shape, batch_dim)); in TorchIndexSelect()
288 to_concat.push_back(Reshape(index, index_shape.dimensions())); in TorchIndexSelect()
289 index = ConcatInDim(builder, to_concat, gather_dnums.index_vector_dim()); in TorchIndexSelect()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | where_op.cc | 74 std::vector<xla::XlaOp> to_concat; in Compile() local
77 to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1})); in Compile()
80 xla::XlaOp result = xla::ConcatInDim(ctx->builder(), to_concat, 1); in Compile()
|
| D | dynamic_stitch_op.cc | 178 std::vector<xla::XlaOp> to_concat(number_of_indices); in Compile() local
186 to_concat[index_num] = in Compile()
190 ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), to_concat, 0)); in Compile()
|
| /external/tensorflow/tensorflow/core/framework/ |
| D | tensor_util_test.cc | 217 std::vector<Tensor> to_concat; in TEST() local
228 to_concat.push_back(tensor); in TEST()
234 TF_ASSERT_OK(tensor::Concat(to_concat, &concated)); in TEST()
|