| 1536 | } |
| 1537 | |
| 1538 | Status GetSparseTensorMetadata(const Buffer& metadata, std::shared_ptr<DataType>* type, |
| 1539 | std::vector<int64_t>* shape, |
| 1540 | std::vector<std::string>* dim_names, |
| 1541 | int64_t* non_zero_length, |
| 1542 | SparseTensorFormat::type* sparse_tensor_format_id) { |
| 1543 | const flatbuf::Message* message = nullptr; |
| 1544 | RETURN_NOT_OK(internal::VerifyMessage(metadata.data(), metadata.size(), &message)); |
| 1545 | auto sparse_tensor = message->header_as_SparseTensor(); |
| 1546 | if (sparse_tensor == nullptr) { |
| 1547 | return Status::IOError( |
| 1548 | "Header-type of flatbuffer-encoded Message is not SparseTensor."); |
| 1549 | } |
| 1550 | int ndim = static_cast<int>(sparse_tensor->shape()->size()); |
| 1551 | |
| 1552 | if (shape || dim_names) { |
| 1553 | for (int i = 0; i < ndim; ++i) { |
| 1554 | auto dim = sparse_tensor->shape()->Get(i); |
| 1555 | |
| 1556 | if (shape) { |
| 1557 | shape->push_back(dim->size()); |
| 1558 | } |
| 1559 | |
| 1560 | if (dim_names) { |
| 1561 | dim_names->push_back(StringFromFlatbuffers(dim->name())); |
| 1562 | } |
| 1563 | } |
| 1564 | } |
| 1565 | |
| 1566 | if (non_zero_length) { |
| 1567 | *non_zero_length = sparse_tensor->non_zero_length(); |
| 1568 | } |
| 1569 | |
| 1570 | if (sparse_tensor_format_id) { |
| 1571 | switch (sparse_tensor->sparseIndex_type()) { |
| 1572 | case flatbuf::SparseTensorIndex::SparseTensorIndex_SparseTensorIndexCOO: |
| 1573 | *sparse_tensor_format_id = SparseTensorFormat::COO; |
| 1574 | break; |
| 1575 | |
| 1576 | case flatbuf::SparseTensorIndex::SparseTensorIndex_SparseMatrixIndexCSX: { |
| 1577 | auto cs = sparse_tensor->sparseIndex_as_SparseMatrixIndexCSX(); |
| 1578 | switch (cs->compressedAxis()) { |
| 1579 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Row: |
| 1580 | *sparse_tensor_format_id = SparseTensorFormat::CSR; |
| 1581 | break; |
| 1582 | |
| 1583 | case flatbuf::SparseMatrixCompressedAxis::SparseMatrixCompressedAxis_Column: |
| 1584 | *sparse_tensor_format_id = SparseTensorFormat::CSC; |
| 1585 | break; |
| 1586 | |
| 1587 | default: |
| 1588 | return Status::Invalid("Invalid value of SparseMatrixCompressedAxis"); |
| 1589 | } |
| 1590 | } break; |
| 1591 | |
| 1592 | case flatbuf::SparseTensorIndex::SparseTensorIndex_SparseTensorIndexCSF: |
| 1593 | *sparse_tensor_format_id = SparseTensorFormat::CSF; |
| 1594 | break; |
| 1595 |
no test coverage detected