aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

transpose n dimensions cpu

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-02 21:14:26 -03:00
parent b7cf03828e
commit a45b930075
6 changed files with 39 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

32
norch/csrc/cpu.cpp
View file

@ -116,7 +116,7 @@ void zeros_like_tensor_cpu(Tensor* tensor, float* result_data) {
}
}
void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
/*void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
int rows = tensor->shape[0];
int cols = tensor->shape[1];
@ -125,6 +125,36 @@ void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
result_data[j * rows + i] = tensor->data[i * cols + j];
}
}
}*/
void transpose_tensor_cpu(Tensor* tensor, float* result_data) {
int* shape = tensor->shape;
int ndim = tensor->ndim;
int* strides = (int*)malloc(ndim * sizeof(int));
int* indices = (int*)calloc(ndim, sizeof(int));
strides[ndim - 1] = 1;
for (int i = ndim - 2; i >= 0; i--) {
strides[i] = strides[i + 1] * shape[i + 1];
}
int idx_result;
for (int idx_source = 0; idx_source < tensor->size; idx_source++) {
idx_result = 0;
for (int dim = 0; dim < ndim; dim++) {
idx_result += indices[dim] * strides[dim];
}
result_data[idx_result] = tensor->data[idx_source];
// Update indices
indices[ndim - 1]++;
for (int dim = ndim - 1; dim > 0; dim--) {
if (indices[dim] == shape[dim]) {
indices[dim] = 0;
indices[dim - 1]++;
}
}
}
}
void assign_tensor_cpu(Tensor* tensor, float* result_data) {