| /external/tensorflow/tensorflow/core/kernels/ |
| D | segment_reduction_ops.cc | 54 const Tensor& segment_ids) { in SegmentReductionValidationHelper() argument
55 OP_REQUIRES(context, TensorShapeUtils::IsVector(segment_ids.shape()), in SegmentReductionValidationHelper()
57 const int64 num_indices = segment_ids.NumElements(); in SegmentReductionValidationHelper()
66 const Tensor& segment_ids) { in SegmentReductionDoValidation() argument
67 SegmentReductionValidationHelper(c, input, segment_ids); in SegmentReductionDoValidation()
82 const Tensor& segment_ids = context->input(1); in Compute() local
84 if (!SegmentReductionDoValidation(context, input, segment_ids)) { in Compute()
88 const int64 num_indices = segment_ids.NumElements(); in Compute()
92 const auto segment_vec = segment_ids.vec<Index>(); in Compute()
220 const Tensor& segment_ids = context->input(1); in ComputeAsync() local
[all …]
|
| D | segment_reduction_ops_gpu.cu.cc | 54 const Index* segment_ids, in SortedSegmentSumCustomKernel() argument
63 Index first_segment_id = segment_ids[input_outer_dim_index_base]; in SortedSegmentSumCustomKernel()
71 segment_ids[input_outer_dim_index_base + j]; in SortedSegmentSumCustomKernel()
108 const Index* segment_ids, in UnsortedSegmentCustomKernel() argument
115 const Index output_segment_index = segment_ids[input_segment_index]; in UnsortedSegmentCustomKernel()
131 typename TTypes<Index>::ConstFlat segment_ids, const Index data_size, in operator ()() argument
150 const Index input_outer_dim_size = segment_ids.dimension(0); in operator ()()
165 segment_ids.data(), data, output.data(), total_stripe_count); in operator ()()
173 typename TTypes<Index>::ConstFlat segment_ids, in operator ()()
192 const Index input_outer_dim_size = segment_ids.dimension(0); in operator ()()
[all …]
|
| D | segment_reduction_ops.h | 45 typename TTypes<Index>::ConstFlat segment_ids,
57 typename TTypes<Index>::ConstFlat segment_ids,
|
| /external/tensorflow/tensorflow/contrib/layers/python/layers/ |
| D | embedding_ops.py | 423 segment_ids = sparse_values.indices[:, 0]
424 if segment_ids.dtype != dtypes.int32:
425 segment_ids = math_ops.cast(segment_ids, dtypes.int32)
434 embeddings = math_ops.sparse_segment_sum(embeddings, idx, segment_ids,
437 embeddings = math_ops.sparse_segment_mean(embeddings, idx, segment_ids,
440 embeddings = math_ops.sparse_segment_sqrt_n(embeddings, idx, segment_ids,
540 segment_ids = sp_values.indices[:, 0]
544 sampled_candidates = array_ops.gather(sampled_candidates, segment_ids)
557 if segment_ids.dtype != dtypes.int32:
558 segment_ids = math_ops.cast(segment_ids, dtypes.int32)
[all …]
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | segment_reduction_ops_test.py | 118 s = tf_op(data=tf_x, segment_ids=indices)
132 math_ops.segment_sum(data=tf_x, segment_ids=indices)
140 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
151 result = math_ops.segment_sum(data=tf_x, segment_ids=indices).eval()
161 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
172 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
181 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
192 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
201 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
213 s = math_ops.segment_sum(data=tf_x, segment_ids=indices)
[all …]
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | embedding_ops.py | 434 segment_ids = sp_ids.indices[:, 0]
435 if segment_ids.dtype != dtypes.int32:
436 segment_ids = math_ops.cast(segment_ids, dtypes.int32)
470 embeddings = math_ops.segment_sum(embeddings, segment_ids, name=name)
472 embeddings = math_ops.segment_sum(embeddings, segment_ids)
473 weight_sum = math_ops.segment_sum(weights, segment_ids)
476 embeddings = math_ops.segment_sum(embeddings, segment_ids)
478 weight_sum = math_ops.segment_sum(weights_squared, segment_ids)
487 embeddings, idx, segment_ids, name=name)
490 embeddings, idx, segment_ids, name=name)
[all …]
|
| D | math_grad_test.py | 160 segment_ids = constant_op.constant([0, 0, 1], dtype=dtypes.int64)
161 segment_min = math_ops.segment_min(data, segment_ids)
169 segment_ids = constant_op.constant([0, 0, 1], dtype=dtypes.int64)
170 segment_max = math_ops.segment_max(data, segment_ids)
179 segment_ids = constant_op.constant([0, 0], dtype=dtypes.int64)
180 segment_min = math_ops.segment_min(data, segment_ids)
189 segment_ids = constant_op.constant([0, 0], dtype=dtypes.int64)
190 segment_max = math_ops.segment_max(data, segment_ids)
|
| D | math_ops.py | 2558 def _unsorted_segment_N(data, segment_ids, num_segments): argument
2563 ones_tensor = array_ops.ones(segment_ids.shape, dtype=data.dtype)
2564 N = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments)
2566 ndims_output = data.shape.ndims - segment_ids.shape.ndims
2573 def unsorted_segment_mean(data, segment_ids, num_segments, name=None): argument
2598 segment_ids = ops.convert_to_tensor(segment_ids)
2599 N = _unsorted_segment_N(data, segment_ids, num_segments)
2600 summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments)
2605 def unsorted_segment_sqrt_n(data, segment_ids, num_segments, name=None): argument
2633 segment_ids = ops.convert_to_tensor(segment_ids)
[all …]
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_UnsortedSegmentSum.pbtxt | 4 name: "segment_ids"
12 Has same shape as data, except for the first `segment_ids.rank`
24 that `segment_ids[j...] == i`. Unlike `SegmentSum`, `segment_ids`
|
| D | api_def_SparseSegmentMean.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
27 Like `SegmentMean`, but `segment_ids` can have rank less than `data`'s first
|
| D | api_def_SparseSegmentMeanWithNumSegments.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
30 Like `SparseSegmentMean`, but allows missing ids in `segment_ids`. If an id is
|
| D | api_def_SparseSegmentSqrtNWithNumSegments.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
32 Like `SparseSegmentSqrtN`, but allows missing ids in `segment_ids`. If an id is
|
| D | api_def_SparseSegmentSumWithNumSegments.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
30 Like `SparseSegmentSum`, but allows missing ids in `segment_ids`. If an id is
|
| D | api_def_SparseSegmentSum.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
27 Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first
|
| D | api_def_SparseSegmentMeanGrad.pbtxt | 16 name: "segment_ids"
18 segment_ids passed to the corresponding SparseSegmentMean op.
|
| D | api_def_SparseSegmentSqrtNGrad.pbtxt | 16 name: "segment_ids"
18 segment_ids passed to the corresponding SparseSegmentSqrtN op.
|
| D | api_def_SparseSegmentSqrtN.pbtxt | 6 A 1-D tensor. Has same rank as `segment_ids`.
10 name: "segment_ids"
|
| D | api_def_SegmentMax.pbtxt | 4 name: "segment_ids"
24 that `segment_ids[j] == i`.
|
| D | api_def_SegmentMin.pbtxt | 4 name: "segment_ids"
24 that `segment_ids[j] == i`.
|
| D | api_def_SegmentProd.pbtxt | 4 name: "segment_ids"
24 that `segment_ids[j] == i`.
|
| D | api_def_UnsortedSegmentProd.pbtxt | 4 name: "segment_ids"
28 that `segment_ids[j] == i`.
|
| D | api_def_SegmentSum.pbtxt | 4 name: "segment_ids"
24 that `segment_ids[j] == i`.
|
| D | api_def_UnsortedSegmentMin.pbtxt | 4 name: "segment_ids"
27 that `segment_ids[j] == i`.
|
| D | api_def_SegmentMean.pbtxt | 4 name: "segment_ids"
24 over `j` such that `segment_ids[j] == i` and `N` is the total number of
|
| D | api_def_UnsortedSegmentMax.pbtxt | 4 name: "segment_ids"
27 that `segment_ids[j] == i`.
|