| 876 | } |
| 877 | |
| 878 | Status WriteStridedTensorData(int dim_index, int64_t offset, int elem_size, |
| 879 | const Tensor& tensor, uint8_t* scratch_space, |
| 880 | io::OutputStream* dst) { |
| 881 | if (dim_index == tensor.ndim() - 1) { |
| 882 | const uint8_t* data_ptr = tensor.raw_data() + offset; |
| 883 | const int64_t stride = tensor.strides()[dim_index]; |
| 884 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 885 | memcpy(scratch_space + i * elem_size, data_ptr, elem_size); |
| 886 | data_ptr += stride; |
| 887 | } |
| 888 | return dst->Write(scratch_space, elem_size * tensor.shape()[dim_index]); |
| 889 | } |
| 890 | for (int64_t i = 0; i < tensor.shape()[dim_index]; ++i) { |
| 891 | RETURN_NOT_OK(WriteStridedTensorData(dim_index + 1, offset, elem_size, tensor, |
| 892 | scratch_space, dst)); |
| 893 | offset += tensor.strides()[dim_index]; |
| 894 | } |
| 895 | return Status::OK(); |
| 896 | } |
| 897 | |
| 898 | Status GetContiguousTensor(const Tensor& tensor, MemoryPool* pool, |
| 899 | std::unique_ptr<Tensor>* out) { |