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

fix sum axis autograd and reversed add broadcasted order

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-17 11:39:26 -03:00
parent 79de915bb2
commit 64fb3f9edb
13 changed files with 58 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

11
norch/csrc/tensor.cpp
View file

@ -149,20 +149,25 @@ extern "C" {
broadcasted_shape[max_ndim - 1 - i] = dim1 > dim2 ? dim1 : dim2;
}
int broadcasted_size = 1;
for (int i = 0; i < max_ndim; i++) {
broadcasted_size *= broadcasted_shape[i];
}
if (strcmp(tensor1->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, tensor1->size * sizeof(float));
cudaMalloc((void **)&result_data, broadcasted_size * sizeof(float));
add_broadcasted_tensor_cuda(tensor1, tensor2, result_data, broadcasted_shape);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
else {
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
float* result_data = (float*)malloc(broadcasted_size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape);
add_broadcasted_tensor_cpu(tensor1, tensor2, result_data, broadcasted_shape, broadcasted_size);
return create_tensor(result_data, broadcasted_shape, max_ndim, tensor1->device);
}
}