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