Searched refs:gather_indices (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | resampler_ops.cc | 138 XlaOp Gather2by2Neighbors(xla::XlaBuilder* b, XlaOp data, XlaOp gather_indices, in Gather2by2Neighbors() argument
158 auto neighbors_data = xla::Gather(data, gather_indices, gather_dim_numbers, in Gather2by2Neighbors()
249 XlaOp gather_indices, XlaOp warp, in CalculateGradData() argument
315 xla::Add(gather_indices, shifting_value, {last_warp_dim}); in CalculateGradData()
360 XlaOp gather_indices, XlaOp data, in CalculateGradWarp() argument
388 xla::Add(gather_indices, shifting_value, {last_warp_dim}); in CalculateGradWarp()
519 auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); in Compile() local
523 ctx->builder(), data, gather_indices, data_channels, warp_shape.dims()); in Compile()
656 auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); in Compile() local
659 ctx, grad_output, ratio, gather_indices, warp, warp_type, warp_shape, in Compile()
[all …]
|
| /external/tensorflow/tensorflow/lite/testing/op_tests/ |
| D | gather_with_constant.py | 45 gather_indices = tf.constant([0, 1], tf.int32)
46 out = tf.equal(reference, tf.gather(gather_input, gather_indices))
|
| /external/tensorflow/tensorflow/compiler/tests/ |
| D | gather_test.py | 161 def _benchmarkGather(self, name, axis, gather_indices, use_xla_jit): argument
169 gather_indices, dtype=dtypes.int32, name='indices')
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | slicing.cc | 214 XlaOp gather_indices = ConcatInDim(builder, to_concat, input_shape.rank()); in TorchGather() local
222 return Gather(input, gather_indices, gather_dnums, slice_sizes); in TorchGather()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | dynamic_padder.cc | 645 HloInstruction* gather_indices = comp->AddInstruction( in RewriteDynamicReshapeCombineInput() local
677 output_shape, reshape_static, gather_indices, gather_dim_numbers, in RewriteDynamicReshapeCombineInput()
|