Fix bug in format converter.
PiperOrigin-RevId: 297738172 Change-Id: I8150a9a9712b4f2118a924699bddff72d1e339f9
This commit is contained in:
parent
a962580295
commit
d63a01cb1f
@ -285,10 +285,12 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
|
||||
}
|
||||
|
||||
const int metadata_idx = 2 * level;
|
||||
const int shape_of_level = dim_metadata_[metadata_idx][0];
|
||||
if (format_[level] == kTfLiteDimDense) {
|
||||
for (int i = 0; i < dim_metadata_[metadata_idx][0]; i++) {
|
||||
for (int i = 0; i < shape_of_level; i++) {
|
||||
indices[level] = i;
|
||||
Populate(src_data, indices, level + 1, i, src_data_ptr);
|
||||
Populate(src_data, indices, level + 1, prev_idx * shape_of_level + i,
|
||||
src_data_ptr);
|
||||
}
|
||||
} else {
|
||||
const auto& array_segments = dim_metadata_[metadata_idx];
|
||||
|
@ -230,6 +230,66 @@ TEST(FormatConverterTest, SimpleTestS1S0) {
|
||||
EXPECT_EQ(data_back, dense_values);
|
||||
}
|
||||
|
||||
TEST(FormatConverterTest, 3DTestS0D1S2) {
|
||||
const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
|
||||
const std::vector<int> dense_shape = {3, 2, 2};
|
||||
const std::vector<int> traversal_order = {0, 1, 2};
|
||||
const std::vector<TfLiteDimensionType> format = {
|
||||
kTfLiteDimSparseCSR, kTfLiteDimDense, kTfLiteDimSparseCSR};
|
||||
FormatConverter<int> converter(dense_shape, traversal_order, format);
|
||||
converter.DenseToSparse(dense_values.data());
|
||||
|
||||
const auto& dim_metadata = converter.GetDimMetadata();
|
||||
const std::vector<int> dm0_0 = {0, 2};
|
||||
const std::vector<int> dm0_1 = {0, 2};
|
||||
const std::vector<int> dm1 = {2};
|
||||
const std::vector<int> dm2_0 = {0, 1, 3, 4, 5};
|
||||
const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
|
||||
|
||||
EXPECT_EQ(dm0_0, dim_metadata[0]);
|
||||
EXPECT_EQ(dm0_1, dim_metadata[1]);
|
||||
EXPECT_EQ(dm1, dim_metadata[2]);
|
||||
EXPECT_EQ(dm2_0, dim_metadata[4]);
|
||||
EXPECT_EQ(dm2_1, dim_metadata[5]);
|
||||
|
||||
const auto& data = converter.GetData();
|
||||
const std::vector<int> expected_data = {6, 9, 8, 5, 7};
|
||||
EXPECT_EQ(expected_data, data);
|
||||
|
||||
converter.SparseToDense(expected_data.data());
|
||||
const auto& data_back = converter.GetData();
|
||||
EXPECT_EQ(data_back, dense_values);
|
||||
}
|
||||
|
||||
TEST(FormatConverterTest, 3DTestD0D1S2) {
|
||||
const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
|
||||
const std::vector<int> dense_shape = {3, 2, 2};
|
||||
const std::vector<int> traversal_order = {0, 1, 2};
|
||||
const std::vector<TfLiteDimensionType> format = {
|
||||
kTfLiteDimDense, kTfLiteDimDense, kTfLiteDimSparseCSR};
|
||||
FormatConverter<int> converter(dense_shape, traversal_order, format);
|
||||
converter.DenseToSparse(dense_values.data());
|
||||
|
||||
const auto& dim_metadata = converter.GetDimMetadata();
|
||||
const std::vector<int> dm0 = {3};
|
||||
const std::vector<int> dm1 = {2};
|
||||
const std::vector<int> dm2_0 = {0, 1, 3, 3, 3, 4, 5};
|
||||
const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
|
||||
|
||||
EXPECT_EQ(dm0, dim_metadata[0]);
|
||||
EXPECT_EQ(dm1, dim_metadata[2]);
|
||||
EXPECT_EQ(dm2_0, dim_metadata[4]);
|
||||
EXPECT_EQ(dm2_1, dim_metadata[5]);
|
||||
|
||||
const auto& data = converter.GetData();
|
||||
const std::vector<int> expected_data = {6, 9, 8, 5, 7};
|
||||
EXPECT_EQ(expected_data, data);
|
||||
|
||||
converter.SparseToDense(expected_data.data());
|
||||
const auto& data_back = converter.GetData();
|
||||
EXPECT_EQ(data_back, dense_values);
|
||||
}
|
||||
|
||||
TEST(FormatConverterTest, 3DTestS0S1S2) {
|
||||
const std::vector<int> dense_values = {1, 7, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0,
|
||||
0, 0, 0, 0, 0, 2, 0, 0, 4, 8, 3, 9};
|
||||
|
Loading…
x
Reference in New Issue
Block a user