Major bug fix in puffernet.h: small but significant discrepency in puffer vs. torch gelu

Author: Joseph Suarez

pufferlib/extensions/puffernet.h

@@ -76,7 +76,7 @@ void _relu(float* input, float* output, int size) {
 
 void _gelu(float* input, float* output, int size) {
     for (int i = 0; i < size; i++) {
-        output[i] = 0.5f*input[i]*(1 + tanhf(0.6628526501011142 * (input[i] + 0.044715f*input[i]*input[i]*input[i])));
+        output[i] = 0.5f*input[i]*(1 + tanhf(0.7978845608028654 * (input[i] + 0.044715f*input[i]*input[i]*input[i])));
     }
 }
 

pufferlib/extensions/puffernet.pyx

@@ -18,6 +18,7 @@ cdef extern from "puffernet.h":
     void _linear(float* input, float* weights, float* bias, float* output,
         int batch_size, int input_dim, int output_dim)
     void _relu(float* input, float* output,int size)
+    void _gelu(float* input, float* output, int size)
     float _sigmoid(float x)
     void _conv2d(float* input, float* weights, float* bias,
         float* output, int batch_size, int in_width, int in_height,
@@ -47,6 +48,9 @@ def puf_linear_layer(cnp.ndarray input, cnp.ndarray weights, cnp.ndarray bias, c
 def puf_relu(cnp.ndarray input, cnp.ndarray output, int size):
     _relu(<float*> input.data, <float*> output.data, size)
 
+def puf_gelu(cnp.ndarray input, cnp.ndarray output, int size):
+    _gelu(<float*> input.data, <float*> output.data, size)
+
 def puf_sigmoid(float x):
     return _sigmoid(x)
 

tests/test_puffernet.py

@@ -42,6 +42,17 @@ def test_puffernet_relu(batch_size=16, input_size=128):
 
     assert_near(input_puffer, output_torch.numpy())
 
+def test_puffernet_gelu(batch_size=16, input_size=128):
+    input_puffer = make_dummy_data(batch_size, input_size)
+
+    input_torch = torch.from_numpy(input_puffer)
+    output_torch = torch.nn.functional.gelu(input_torch, approximate='tanh').detach()
+
+    # PufferNet done second because it is in-place on the input
+    puffernet.puf_gelu(input_puffer, input_puffer, batch_size*input_size)
+
+    assert_near(input_puffer, output_torch.numpy())
+
 def test_puffernet_sigmoid(n=1024, epsilon=1e-4):
     input_np = make_dummy_data(n)
 
@@ -247,9 +258,8 @@ def test_nmmo3(batch_size=1, input_size=512, hidden_size=512):
     pass
 
 if __name__ == '__main__':
-    test_nmmo3()
-    exit()
     test_puffernet_relu()
+    test_puffernet_gelu()
     test_puffernet_sigmoid()
     test_puffernet_linear_layer()
     test_puffernet_convolution_layer()
@@ -260,3 +270,4 @@ def test_nmmo3(batch_size=1, input_size=512, hidden_size=512):
     test_puffernet_one_hot()
     test_puffernet_cat_dim1()
     test_puffernet_argmax_multidiscrete()
+    #test_nmmo3()