fastgrnncuda: adds support for different non-linearities

Author: MJ10

pytorch/edgeml_pytorch/cuda/fastgrnn_cuda.cpp

@@ -4,22 +4,24 @@
 
 std::vector<torch::Tensor> fastgrnn_cuda_forward(
     torch::Tensor input,
-    torch::Tensor w,
-    torch::Tensor u,
+    torch::Tensor W,
+    torch::Tensor U,
     torch::Tensor bias_gate,
     torch::Tensor bias_update,
     torch::Tensor zeta,
     torch::Tensor nu,
-    torch::Tensor old_h);
+    torch::Tensor old_h,
+    int z_non_linearity);
 
 std::vector<torch::Tensor> fastgrnn_cuda_backward(
     torch::Tensor grad_h,
     torch::Tensor input,
     torch::Tensor old_h,
     torch::Tensor zeta,
     torch::Tensor nu,
-    torch::Tensor w,
-    torch::Tensor u,
+    torch::Tensor W,
+    torch::Tensor U,
+    int z_non_linearity,
     torch::Tensor z,
     torch::Tensor h_prime);
 
@@ -29,23 +31,24 @@ std::vector<torch::Tensor> fastgrnn_cuda_backward(
 
 std::vector<torch::Tensor> fastgrnn_forward(
     torch::Tensor input,
-    torch::Tensor w,
-    torch::Tensor u,
+    torch::Tensor W,
+    torch::Tensor U,
     torch::Tensor bias_gate,
     torch::Tensor bias_update,
     torch::Tensor zeta,
     torch::Tensor nu,
-    torch::Tensor old_h) {
+    torch::Tensor old_h,
+    int z_non_linearity) {
   CHECK_INPUT(input);
-  CHECK_INPUT(w);
-  CHECK_INPUT(u);
+  CHECK_INPUT(W);
+  CHECK_INPUT(U);
   CHECK_INPUT(bias_gate);
   CHECK_INPUT(bias_update);
   CHECK_INPUT(zeta);
   CHECK_INPUT(nu);
   CHECK_INPUT(old_h);
 
-  return fastgrnn_cuda_forward(input, w, u, bias_gate, bias_update, zeta, nu, old_h);
+  return fastgrnn_cuda_forward(input, W, U, bias_gate, bias_update, zeta, nu, old_h, z_non_linearity);
 }
 
 std::vector<torch::Tensor> fastgrnn_backward(
@@ -54,21 +57,22 @@ std::vector<torch::Tensor> fastgrnn_backward(
     torch::Tensor old_h,
     torch::Tensor zeta,
     torch::Tensor nu,
-    torch::Tensor w,
-    torch::Tensor u,
+    torch::Tensor W,
+    torch::Tensor U,
     torch::Tensor z,
-    torch::Tensor h_prime) {
+    torch::Tensor h_prime,
+    int z_non_linearity) {
   CHECK_INPUT(grad_h);
   CHECK_INPUT(input);
   CHECK_INPUT(old_h);
   CHECK_INPUT(zeta);
   CHECK_INPUT(nu);
   CHECK_INPUT(z);
   CHECK_INPUT(h_prime);
-  CHECK_INPUT(w);
-  CHECK_INPUT(u);
+  CHECK_INPUT(W);
+  CHECK_INPUT(U);
 
-  return fastgrnn_cuda_backward(grad_h, input, old_h, zeta, nu, w, u, z, h_prime);
+  return fastgrnn_cuda_backward(grad_h, input, old_h, zeta, nu, W, U, z_non_linearity, z, h_prime);
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {