message: Set the type of sparsity metadata accordingly to further reduce the model size.
diff: @@ -1406,22 +1406,67 @@ BufferOffset<tflite::SparsityParameters> Translator::BuildSparsityParameters(
  for (int j = 0; j < segments.size(); j++) {
  vector_segments[j] = segments[j].dyn_cast<mlir::IntegerAttr>().getInt();
  }
  - auto array_segments =
  - tflite::CreateInt32Vector(builder_,
  - builder_.CreateVector(vector_segments))
  + tflite::SparseIndexVector segments_type;
  + BufferOffset<void> array_segments;
  + // The segment array is sorted.
  + // TODO(b/147449640): Clean this up with util functions.
  + int max_of_segments = vector_segments[segments.size() - 1];
  + if (max_of_segments <= UINT8_MAX) {
  + segments_type = tflite::SparseIndexVector_Uint8Vector;
  + std::vector<uint8_t> uint8_vector(vector_segments.begin(),
  + vector_segments.end());
  + array_segments = tflite::CreateUint8Vector(
  + builder_, builder_.CreateVector(uint8_vector))
  .Union();
  + } else if (max_of_segments <= UINT16_MAX) {
  + segments_type = tflite::SparseIndexVector_Uint16Vector;
  + std::vector<uint16_t> uint16_vector(vector_segments.begin(),
  + vector_segments.end());
  + array_segments = tflite::CreateUint16Vector(
  + builder_, builder_.CreateVector(uint16_vector))
  + .Union();
  + } else {
  + segments_type = tflite::SparseIndexVector_Int32Vector;
  + array_segments = tflite::CreateInt32Vector(
  + builder_, builder_.CreateVector(vector_segments))
  + .Union();
  + }
  +
  auto indices = dim_metadata.indices();
  std::vector<int> vector_indices(indices.size(), 0);
  + int max_of_indices = 0;
  for (int j = 0; j < indices.size(); j++) {
  vector_indices[j] = indices[j].dyn_cast<mlir::IntegerAttr>().getInt();
  - }
  - auto array_indices = tflite::CreateInt32Vector(
  + if (vector_indices[j] > max_of_indices) {
  + max_of_indices = vector_indices[j];
  + }
  + }
  + tflite::SparseIndexVector indices_type;
  + BufferOffset<void> array_indices;
  + if (max_of_indices <= UINT8_MAX) {
  + indices_type = tflite::SparseIndexVector_Uint8Vector;
  + std::vector<uint8_t> uint8_vector(vector_indices.begin(),
  + vector_indices.end());
  + array_indices = tflite::CreateUint8Vector(
  + builder_, builder_.CreateVector(uint8_vector))
  + .Union();
  + } else if (max_of_indices <= UINT16_MAX) {
  + indices_type = tflite::SparseIndexVector_Uint16Vector;
  + std::vector<uint16_t> uint16_vector(vector_indices.begin(),
  + vector_indices.end());
  + array_indices = tflite::CreateUint16Vector(
  + builder_, builder_.CreateVector(uint16_vector))
  + .Union();
  + } else {
  + indices_type = tflite::SparseIndexVector_Int32Vector;
  + array_indices = tflite::CreateInt32Vector(
  builder_, builder_.CreateVector(vector_indices))
  .Union();
  + }
  +
  fb_dim_metadata[i] = tflite::CreateDimensionMetadata(
  - builder_, tflite::DimensionType_SPARSE_CSR, 0,
  - tflite::SparseIndexVector_Int32Vector, array_segments,
  - tflite::SparseIndexVector_Int32Vector, array_indices);
  + builder_, tflite::DimensionType_SPARSE_CSR, 0, segments_type,
  + array_segments, indices_type, array_indices);
  }
  }
  