| 544 | |
| 545 | template <typename TYPE> |
| 546 | int64_t StridedTensorCountNonZero(int dim_index, int64_t offset, const Tensor& tensor) { |
| 547 | using c_type = typename TYPE::c_type; |
| 548 | const c_type zero = c_type(0); |
| 549 | int64_t nnz = 0; |
| 550 | if (dim_index == tensor.ndim() - 1) { |
| 551 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 552 | const auto* ptr = tensor.raw_data() + offset + i * tensor.strides()[dim_index]; |
| 553 | auto& elem = *reinterpret_cast<const c_type*>(ptr); |
| 554 | if (elem != zero) ++nnz; |
| 555 | } |
| 556 | return nnz; |
| 557 | } |
| 558 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 559 | nnz += StridedTensorCountNonZero<TYPE>(dim_index + 1, offset, tensor); |
| 560 | offset += tensor.strides()[dim_index]; |
| 561 | } |
| 562 | return nnz; |
| 563 | } |
| 564 | |
| 565 | template <typename TYPE> |
| 566 | int64_t ContiguousTensorCountNonZero(const Tensor& tensor) { |