MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / PopulateInternal

Method PopulateInternal

tensorflow/compiler/xla/literal.h:990–1032  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

988
989template <typename NativeT, typename FnType>
990Status 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}
1033template <typename NativeT, typename FnType>
1034Status MutableLiteralBase::Populate(const FnType& generator) {
1035 return PopulateInternal<NativeT>(generator, /*parallel=*/false);

Callers

nothing calls this directly

Calls 11

ForEachIndexFunction · 0.85
init_functionFunction · 0.85
AsInt64SliceFunction · 0.70
shapeFunction · 0.50
copyFunction · 0.50
rankMethod · 0.45
element_typeMethod · 0.45
dimensionsMethod · 0.45
beginMethod · 0.45
endMethod · 0.45
atMethod · 0.45

Tested by

no test coverage detected