Spaces:
Sleeping
Sleeping
File size: 1,814 Bytes
899c526 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
#ifndef LIETORCH_GPU_H_
#define LIETORCH_GPU_H_
#include <vector>
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
// unary operations
torch::Tensor exp_forward_gpu(int, torch::Tensor);
std::vector<torch::Tensor> exp_backward_gpu(int, torch::Tensor, torch::Tensor);
torch::Tensor log_forward_gpu(int, torch::Tensor);
std::vector<torch::Tensor> log_backward_gpu(int, torch::Tensor, torch::Tensor);
torch::Tensor inv_forward_gpu(int, torch::Tensor);
std::vector<torch::Tensor> inv_backward_gpu(int, torch::Tensor, torch::Tensor);
// binary operations
torch::Tensor mul_forward_gpu(int, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> mul_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
torch::Tensor adj_forward_gpu(int, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> adj_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
torch::Tensor adjT_forward_gpu(int, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> adjT_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
torch::Tensor act_forward_gpu(int, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> act_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
torch::Tensor act4_forward_gpu(int, torch::Tensor, torch::Tensor);
std::vector<torch::Tensor> act4_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
// conversion operations
// std::vector<torch::Tensor> to_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
// std::vector<torch::Tensor> from_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
// utility operators
torch::Tensor orthogonal_projector_gpu(int, torch::Tensor);
torch::Tensor as_matrix_forward_gpu(int, torch::Tensor);
torch::Tensor jleft_forward_gpu(int, torch::Tensor, torch::Tensor);
#endif
|