Searched refs:gather_output (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | gather_op.cc | 39 xla::XlaOp* gather_output) { in XlaGather() argument
84 *gather_output = in XlaGather()
151 *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes); in XlaGather()
158 int batch_dims, xla::XlaOp* gather_output) { in XlaGatherWithBatchDimsOpImpl() argument
223 *gather_output = xla::TorchIndexSelect(input, indices, *axis, batch_dims); in XlaGatherWithBatchDimsOpImpl()
230 index_type, context->builder(), gather_output)); in XlaGatherWithBatchDimsOpImpl()
|
| D | gather_op_helpers.h | 40 xla::XlaOp* gather_output);
48 int batch_dims, xla::XlaOp* gather_output);
|
| /external/tensorflow/tensorflow/compiler/aot/tests/ |
| D | test_graph_tfgather.config.pbtxt | 15 id { node_name: "gather_output" }
|