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