Searched refs:TorchIndexSelect (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | slicing_test.cc | 164 TorchIndexSelect(input, index, 0); in XLA_TEST_F()
184 TorchIndexSelect(input, index, 1); in XLA_TEST_F()
198 TorchIndexSelect(input, index, 1); in XLA_TEST_F()
215 TorchIndexSelect(input, index, 0); in XLA_TEST_F()
229 TorchIndexSelect(input, index, 0); in XLA_TEST_F()
245 TorchIndexSelect(input, index, 1, 1); in XLA_TEST_F()
|
| D | slicing.h | 83 XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim,
|
| D | slicing.cc | 169 return TorchIndexSelect(input, index, 0); in TorchGather()
263 XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64 dim, int64 batch_dims) { in TorchIndexSelect() function
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | data_format_ops.cc | 73 xla::TorchIndexSelect(dst_indices, src_indices, /*dim=*/0); in Compile()
154 xla::XlaOp output = xla::TorchIndexSelect(ctx->Input(0), indices, 0); in Compile()
|
| D | gather_op.cc | 223 *gather_output = xla::TorchIndexSelect(input, indices, *axis, batch_dims); in XlaGatherWithBatchDimsOpImpl()
|