Unit tests and some stubs:

tests/nn/data_parallel/test_data_parallel.py

@@ -59,7 +59,7 @@ def run_parallelize_a_transformers_and_inference(
     assert torch.allclose(outputs["loss"], REF_LOSS)
 
 
-def test_data_parllel_fused_bias_gelu_bias_dropout_fwd():
+def test_data_parallel_fused_bias_gelu_bias_dropout_fwd():
     # TODO
     pass
 

tests/nn/tensor_parallel/test_tensor_parallel.py

@@ -94,7 +94,7 @@ def get_leaf_modules(model):
     assert torch.allclose(generated_tokens, REF_GENERATED_TOKENS)
 
 
-def test_data_parllel_fused_bias_gelu_bias_dropout_fwd():
+def test_tensor_parallel_fused_bias_gelu_bias_dropout_fwd():
     # TODO
     pass
 

tests/nn/test_fusion.py

@@ -0,0 +1,28 @@
+import pytest
+import torch
+from torch.nn import Dropout, GELU
+
+from pipegoose.nn.fusion import FusedBiasDropout, FusedBiasGelu
+
+
+def test_FusedBiasDropout():
+    dropout = FusedBiasDropout()
+    input = torch.randn(20, 16)
+    args = (0.5, False)
+    expected = Dropout(*args)(input)
+    actual = dropout(input)
+
+    assert actual.size() == expected.size()
+    assert torch.allclose(actual, expected)
+    assert FusedBiasDropout(*args).represents == Dropout
+
+
+def test_FusedBiasGelu():
+    gelu = FusedBiasGelu()
+    input = torch.randn(20, 16)
+    expected = GELU()(input)
+    actual = gelu(input)
+
+    assert actual.size() == expected.size()
+    assert torch.allclose(actual, expected)
+    assert FusedBiasGelu().represents == GELU