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