| 988 | |
| 989 | template <typename NativeT, typename FnType> |
| 990 | Status MutableLiteralBase::PopulateInternal(const FnType& generator, |
| 991 | bool parallel) { |
| 992 | const Shape& this_shape = shape(); |
| 993 | const int64 rank = this_shape.rank(); |
| 994 | TF_RET_CHECK(LayoutUtil::IsDenseArray(this_shape)); |
| 995 | TF_RET_CHECK(this_shape.element_type() == |
| 996 | primitive_util::NativeToPrimitiveType<NativeT>()); |
| 997 | absl::Span<NativeT> literal_data = data<NativeT>(); |
| 998 | if (rank > 0) { |
| 999 | StrideConfig stride_config(this_shape, this_shape, |
| 1000 | AsInt64Slice(this_shape.dimensions())); |
| 1001 | int64 minor_dimension_size = |
| 1002 | ShapeUtil::GetDimension(this_shape, stride_config.minor_dimension); |
| 1003 | |
| 1004 | auto init_function = [&](absl::Span<const int64> indexes) { |
| 1005 | DimensionVector minor_scan_indexes(rank, 0); |
| 1006 | const int64 index = |
| 1007 | IndexUtil::MultidimensionalIndexToLinearIndex(shape(), indexes); |
| 1008 | std::copy(indexes.begin(), indexes.end(), minor_scan_indexes.begin()); |
| 1009 | for (int64 i = 0; i < minor_dimension_size; ++i) { |
| 1010 | minor_scan_indexes[stride_config.minor_dimension] = i; |
| 1011 | literal_data.at(index + i) = generator(minor_scan_indexes); |
| 1012 | } |
| 1013 | }; |
| 1014 | if (parallel) { |
| 1015 | ShapeUtil::ForEachIndexParallel(this_shape, stride_config.base, |
| 1016 | stride_config.dimensions, |
| 1017 | stride_config.step, init_function); |
| 1018 | } else { |
| 1019 | ShapeUtil::ForEachIndex( |
| 1020 | this_shape, stride_config.base, stride_config.dimensions, |
| 1021 | stride_config.step, |
| 1022 | [&init_function](absl::Span<const int64> indexes) { |
| 1023 | init_function(indexes); |
| 1024 | return true; |
| 1025 | }); |
| 1026 | } |
| 1027 | } else { |
| 1028 | // For scalars. |
| 1029 | literal_data.at(0) = generator({}); |
| 1030 | } |
| 1031 | return Status::OK(); |
| 1032 | } |
| 1033 | template <typename NativeT, typename FnType> |
| 1034 | Status MutableLiteralBase::Populate(const FnType& generator) { |
| 1035 | return PopulateInternal<NativeT>(generator, /*parallel=*/false); |
nothing calls this directly
no test coverage detected