experiment with binary dropout

Author: Z80coder

neurallogic/neural_logic_net.py

@@ -1,29 +1,32 @@
 from enum import Enum
 from flax import linen as nn
 
-NetType = Enum('NetType', ['Soft', 'Hard', 'Symbolic'])
+NetType = Enum("NetType", ["Soft", "Hard", "Symbolic"])
+
 
 def select(soft, hard, symbolic):
     def selector(type: NetType):
-        return {
-            NetType.Soft: soft,
-            NetType.Hard: hard,
-            NetType.Symbolic: symbolic
-        }[type]
+        return {NetType.Soft: soft, NetType.Hard: hard, NetType.Symbolic: symbolic}[
+            type
+        ]
+
     return selector
 
+
 def net(f):
     class SoftNet(nn.Module):
         @nn.compact
-        def __call__(self, x):
-            return f(NetType.Soft, x)
-    class HardNet(nn.Module): 
+        def __call__(self, x, **kwargs):
+            return f(NetType.Soft, x, **kwargs)
+
+    class HardNet(nn.Module):
         @nn.compact
         def __call__(self, x):
             return f(NetType.Hard, x)
+
     class SymbolicNet(nn.Module):
         @nn.compact
         def __call__(self, x):
             return f(NetType.Symbolic, x)
-    return SoftNet(), HardNet(), SymbolicNet()
 
+    return SoftNet(), HardNet(), SymbolicNet()

tests/test_mnist.py

@@ -325,5 +325,5 @@ def test_mnist():
     )
 
     # Check symbolic net
-    # _, hard, symbolic = neural_logic_net.net(lambda type, x: nln(type, x))
-    # check_symbolic((soft, hard, symbolic), (train_ds, test_ds), trained_state)
+    _, hard, symbolic = neural_logic_net.net(lambda type, x: nln(type, x))
+    check_symbolic((soft, hard, symbolic), (train_ds, test_ds), trained_state)