Home
last modified time | relevance | path

Searched refs:gather_indices (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dresampler_ops.cc138 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/contrib/seq2seq/python/ops/
Dbeam_search_decoder.py1027 gather_indices=word_indices,
1048 gather_indices=next_beam_ids,
1066 gather_indices=next_beam_ids,
1075 gather_indices=next_beam_ids,
1089 gather_indices=next_beam_ids,
1305 def _maybe_tensor_gather_helper(gather_indices, gather_from, batch_size, argument
1334 gather_indices=gather_indices,
1343 def _tensor_gather_helper(gather_indices, argument
1374 gather_indices = array_ops.reshape(gather_indices + range_, [-1])
1376 array_ops.reshape(gather_from, gather_shape), gather_indices)
/external/tensorflow/tensorflow/contrib/factorization/python/ops/
Dfactorization_ops.py885 gather_indices = update_row_indices
891 gather_indices = update_col_indices
896 right_factors, gather_indices, partition_strategy="div")
937 col_wt, gather_indices, partition_strategy="div")
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dslicing.cc152 XlaOp gather_indices = ConcatInDim(builder, to_concat, input_shape.rank()); in TorchGather() local
160 return Gather(input, gather_indices, gather_dnums, slice_sizes); in TorchGather()
/external/tensorflow/tensorflow/compiler/tests/
Dgather_test.py161 def _benchmarkGather(self, name, axis, gather_indices, use_xla_jit): argument
169 gather_indices, dtype=dtypes.int32, name='indices')
/external/tensorflow/tensorflow/contrib/tensor_forest/python/
Dtensor_forest.py458 gather_indices = array_ops.squeeze(
462 tree_data = array_ops.gather(processed_dense_features, gather_indices)
463 tree_labels = array_ops.gather(labels, gather_indices)
/external/tensorflow/tensorflow/lite/testing/
Dgenerate_examples.py1494 gather_indices = tf.constant([0, 1], tf.int32)
1495 out = tf.equal(reference, tf.gather(gather_input, gather_indices))