| 484 | |
| 485 | template <typename TYPE> |
| 486 | int64_t StridedTensorCountNonZero(int dim_index, int64_t offset, const Tensor& tensor) { |
| 487 | using c_type = typename TYPE::c_type; |
| 488 | const c_type zero = c_type(0); |
| 489 | int64_t nnz = 0; |
| 490 | if (dim_index == tensor.ndim() - 1) { |
| 491 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 492 | const auto* ptr = tensor.raw_data() + offset + i * tensor.strides()[dim_index]; |
| 493 | auto& elem = *reinterpret_cast<const c_type*>(ptr); |
| 494 | if (elem != zero) ++nnz; |
| 495 | } |
| 496 | return nnz; |
| 497 | } |
| 498 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 499 | nnz += StridedTensorCountNonZero<TYPE>(dim_index + 1, offset, tensor); |
| 500 | offset += tensor.strides()[dim_index]; |
| 501 | } |
| 502 | return nnz; |
| 503 | } |
| 504 | |
| 505 | template <typename TYPE> |
| 506 | int64_t ContiguousTensorCountNonZero(const Tensor& tensor) { |