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

autograd add function

Browse source
This commit is contained in:
lucasdelimanogueira 2024-04-30 20:09:08 -03:00
parent d4751341bd
commit 407242bd7e
15 changed files with 229 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

76
norch/csrc/tensor.cpp
View file

@ -461,4 +461,78 @@ extern "C" {
stride *= new_shape[i];
}
}
}
}
Tensor* ones_like_tensor(Tensor* tensor) {
char* device = (char*)malloc(strlen(tensor->device) + 1);
if (device != NULL) {
strcpy(device, tensor->device);
} else {
fprintf(stderr, "Memory allocation failed\n");
exit(-1);
}
int ndim = tensor->ndim;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
for (int i = 0; i < ndim; i++) {
shape[i] = tensor->shape[i];
}
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
ones_like_tensor_cuda(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {
float* result_data = (float*)malloc(tensor->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
ones_like_tensor_cpu(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
}
Tensor* zeros_like_tensor(Tensor* tensor) {
char* device = (char*)malloc(strlen(tensor->device) + 1);
if (device != NULL) {
strcpy(device, tensor->device);
} else {
fprintf(stderr, "Memory allocation failed\n");
exit(-1);
}
int ndim = tensor->ndim;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
for (int i = 0; i < ndim; i++) {
shape[i] = tensor->shape[i];
}
if (strcmp(tensor->device, "cuda") == 0) {
float* result_data;
cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
zeros_like_tensor_cuda(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
else {
float* result_data = (float*)malloc(tensor->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
zeros_like_tensor_cpu(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
}