| 1109 | |
| 1110 | template <typename SparseIndexType> |
| 1111 | Status MakeSparseMatrixIndexCSX(FBB& fbb, const SparseIndexType& sparse_index, |
| 1112 | const std::vector<BufferMetadata>& buffers, |
| 1113 | flatbuf::SparseTensorIndex* fb_sparse_index_type, |
| 1114 | Offset* fb_sparse_index, size_t* num_buffers) { |
| 1115 | *fb_sparse_index_type = |
| 1116 | flatbuf::SparseTensorIndex::SparseTensorIndex_SparseMatrixIndexCSX; |
| 1117 | |
| 1118 | // We assume that the value type of indptr tensor is an integer. |
| 1119 | const auto& indptr_value_type = |
| 1120 | checked_cast<const IntegerType&>(*sparse_index.indptr()->type()); |
| 1121 | auto indptr_type_offset = flatbuf::CreateInt(fbb, indptr_value_type.bit_width(), |
| 1122 | indptr_value_type.is_signed()); |
| 1123 | |
| 1124 | const BufferMetadata& indptr_metadata = buffers[0]; |
| 1125 | flatbuf::Buffer indptr(indptr_metadata.offset, indptr_metadata.length); |
| 1126 | |
| 1127 | // We assume that the value type of indices tensor is an integer. |
| 1128 | const auto& indices_value_type = |
| 1129 | checked_cast<const IntegerType&>(*sparse_index.indices()->type()); |
| 1130 | auto indices_type_offset = flatbuf::CreateInt(fbb, indices_value_type.bit_width(), |
| 1131 | indices_value_type.is_signed()); |
| 1132 | |
| 1133 | const BufferMetadata& indices_metadata = buffers[1]; |
| 1134 | flatbuf::Buffer indices(indices_metadata.offset, indices_metadata.length); |
| 1135 | |
| 1136 | auto compressedAxis = SparseMatrixCompressedAxis<SparseIndexType>::value; |
| 1137 | *fb_sparse_index = |
| 1138 | flatbuf::CreateSparseMatrixIndexCSX(fbb, compressedAxis, indptr_type_offset, |
| 1139 | &indptr, indices_type_offset, &indices) |
| 1140 | .Union(); |
| 1141 | *num_buffers = 2; |
| 1142 | return Status::OK(); |
| 1143 | } |
| 1144 | |
| 1145 | Status MakeSparseTensorIndexCSF(FBB& fbb, const SparseCSFIndex& sparse_index, |
| 1146 | const std::vector<BufferMetadata>& buffers, |
no test coverage detected