Searched refs:segment_ids_shape (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | segment_reduction_ops_gpu.cu.cc | 137 const TensorShape& segment_ids_shape, in operator ()() argument
150 if (data_size == 0 || segment_ids_shape.num_elements() == 0) { in operator ()()
184 void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape, in operator ()()
198 if (data_size == 0 || segment_ids_shape.num_elements() == 0) { in operator ()()
|
| D | segment_reduction_ops.h | 53 const Index output_rows, const TensorShape& segment_ids_shape,
64 void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
|
| D | segment_reduction_ops_impl.h | 320 void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
338 "segment_ids", SliceDebugString(segment_ids_shape, i),
|
| /external/tensorflow/tensorflow/core/ops/ |
| D | math_ops.cc | 1124 ShapeHandle segment_ids_shape; in SegmentReductionShapeFn() local
1126 TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &segment_ids_shape)); in SegmentReductionShapeFn()
1145 ShapeHandle segment_ids_shape; in SparseSegmentReductionShapeFn() local
1146 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &segment_ids_shape)); in SparseSegmentReductionShapeFn()
1150 TF_RETURN_IF_ERROR(c->Merge(indices_shape, segment_ids_shape, &unused)); in SparseSegmentReductionShapeFn()
1207 ShapeHandle segment_ids_shape; in SparseSegmentReductionWithNumSegmentsShapeFn() local
1208 TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &segment_ids_shape)); in SparseSegmentReductionWithNumSegmentsShapeFn()
1215 TF_RETURN_IF_ERROR(c->Merge(indices_shape, segment_ids_shape, &unused)); in SparseSegmentReductionWithNumSegmentsShapeFn()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | math_ops.py | 4184 segment_ids_shape = array_ops.shape_internal(segment_ids)
4185 ones_tensor = array_ops.ones(segment_ids_shape, dtype=data.dtype)
|
| /external/tensorflow/tensorflow/lite/kernels/internal/reference/ |
| D | reference_ops.h | 1800 const RuntimeShape& segment_ids_shape, in SegmentSum() argument
|