cleanup

Author: Z80coder

neurallogic/hard_not.py

@@ -28,9 +28,9 @@ def hard_not(w: bool, x: bool) -> bool:
 hard_not_neuron = jax.vmap(hard_not, 0, 0)
 
 
-soft_not_layer = jax.vmap(soft_not_neuron, (0, 0), 0)
+soft_not_layer = jax.vmap(soft_not_neuron, (0, None), 0)
 
-hard_not_layer = jax.vmap(hard_not_neuron, (0, 0), 0)
+hard_not_layer = jax.vmap(hard_not_neuron, (0, None), 0)
 
 
 class SoftNotLayer(nn.Module):

neurallogic/harden.py

@@ -11,41 +11,58 @@ def harden_float(x: float) -> bool:
 
 harden_array = jax.vmap(harden_float, 0, 0)
 
+
 @dispatch
 def harden(x: float):
     if numpy.isnan(x):
         return x
     return harden_float(x)
 
+
 @dispatch
 def harden(x: list):
     return symbolic_primitives.map_at_elements(x, harden_float)
 
+
 @dispatch
 def harden(x: numpy.ndarray):
     return harden_array(x)
 
+
 @dispatch
 def harden(x: jax.numpy.ndarray):
     return harden_array(x)
 
+
 @dispatch
 def harden(x: dict):
-    return symbolic_primitives.map_at_elements(x, harden_float)
+    return symbolic_primitives.map_at_elements(x, harden)
+
 
 @dispatch
 def harden(x: flax.core.FrozenDict):
-    return flax.core.FrozenDict(symbolic_primitives.map_at_elements(x.unfreeze(), harden_float))
+    return flax.core.FrozenDict(
+        symbolic_primitives.map_at_elements(x.unfreeze(), harden)
+    )
+
 
 @dispatch
 def harden(*args):
     if len(args) == 1:
         return harden(args[0])
     return tuple([harden(arg) for arg in args])
 
+
 @dispatch
 def map_keys_nested(f, d: dict) -> dict:
-    return {f(k): map_keys_nested(f, v) if isinstance(v, dict) else v for k, v in d.items()}
+    return {
+        f(k): map_keys_nested(f, v) if isinstance(v, dict) else v for k, v in d.items()
+    }
+
 
 def hard_weights(weights):
-    return flax.core.FrozenDict(map_keys_nested(lambda str: str.replace("Soft", "Hard"), harden(weights.unfreeze())))
+    return flax.core.FrozenDict(
+        map_keys_nested(
+            lambda str: str.replace("Soft", "Hard"), harden(weights.unfreeze())
+        )
+    )

neurallogic/real_encoder.py

@@ -1,4 +1,9 @@
+from typing import Callable
+
 import jax
+from flax import linen as nn
+
+from neurallogic import neural_logic_net, symbolic_generation
 
 
 def soft_real_encoder(t: float, x: float) -> float:
@@ -29,3 +34,51 @@ def hard_real_encoder(t: float, x: float) -> bool:
 hard_real_encoder_layer = jax.vmap(hard_real_encoder_neuron, (0, 0), 0)
 
 
+class SoftRealEncoderLayer(nn.Module):
+    bits_per_real: int
+    thresholds_init: Callable = nn.initializers.uniform(1.0)
+    dtype: jax.numpy.dtype = jax.numpy.float32
+
+    @nn.compact
+    def __call__(self, x):
+        thresholds_shape = (jax.numpy.shape(x)[-1], self.bits_per_real)
+        thresholds = self.param("thresholds", self.thresholds_init, thresholds_shape, self.dtype)
+        x = jax.numpy.asarray(x, self.dtype)
+        return soft_real_encoder_layer(thresholds, x)
+
+
+class HardRealEncoderLayer(nn.Module):
+    bits_per_real: int
+
+    @nn.compact
+    def __call__(self, x):
+        thresholds_shape = (jax.numpy.shape(x)[-1], self.bits_per_real)
+        thresholds = self.param("thresholds", nn.initializers.constant(0.0), thresholds_shape)
+        return hard_real_encoder_layer(thresholds, x)
+
+
+class SymbolicRealEncoderLayer:
+    def __init__(self, bits_per_real):
+        self.bits_per_real = bits_per_real
+        self.hard_real_encoder_layer = HardRealEncoderLayer(self.bits_per_real)
+
+    def __call__(self, x):
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
+            self.hard_real_encoder_layer, x
+        )
+        return symbolic_generation.symbolic_expression(jaxpr, x)
+
+
+real_encoder_layer = neural_logic_net.select(
+    lambda bits_per_real, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SoftRealEncoderLayer(
+        bits_per_real, weights_init, dtype
+    ),
+    lambda bits_per_real, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: HardRealEncoderLayer(bits_per_real),
+    lambda bits_per_real, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SymbolicRealEncoderLayer(bits_per_real),
+)

tests/test_mnist.py

@@ -195,7 +195,7 @@ def get_config():
     # Always commit with num_epochs = 1 for short test time
     config.momentum = 0.9
     config.batch_size = 128
-    config.num_epochs = 1000
+    config.num_epochs = 2
     return config
 
 

tests/test_real_encoder.py

@@ -1,32 +1,32 @@
 from typing import Callable
 import numpy
 import jax
+from jax import random
 
-from neurallogic import harden, real_encoder, symbolic_generation
+from neurallogic import harden, real_encoder, symbolic_generation, neural_logic_net
 
 
 def check_consistency(soft: Callable, hard: Callable, expected, *args):
-    print(f'\nchecking consistency for {soft.__name__}')
+    # print(f'\nchecking consistency for {soft.__name__}')
     # Check that the soft function performs as expected
     soft_output = soft(*args)
-    print(f'Expected: {expected}, Actual soft_output: {soft_output}')
+    # print(f'Expected: {expected}, Actual soft_output: {soft_output}')
     assert numpy.allclose(soft_output, expected, equal_nan=True)
 
     # Check that the hard function performs as expected
     # N.B. We don't harden the inputs because the hard_bit expects real-valued inputs
     hard_expected = harden.harden(expected)
-    hard_output = hard(*args) 
-    print(f'Expected: {hard_expected}, Actual hard_output: {hard_output}')
+    hard_output = hard(*args)
+    # print(f'Expected: {hard_expected}, Actual hard_output: {hard_output}')
     assert numpy.allclose(hard_output, hard_expected, equal_nan=True)
 
     # Check that the jaxpr performs as expected
     symbolic_f = symbolic_generation.make_symbolic_jaxpr(hard, *args)
     symbolic_output = symbolic_generation.eval_symbolic(symbolic_f, *args)
-    print(f'Expected: {hard_expected}, Actual symbolic_output: {symbolic_output}')
+    # print(f'Expected: {hard_expected}, Actual symbolic_output: {symbolic_output}')
     assert numpy.allclose(symbolic_output, hard_expected, equal_nan=True)
 
 
-
 def test_activation():
     test_data = [
         [[1.0, 1.0], 0.5],
@@ -40,7 +40,11 @@ def test_activation():
     ]
     for input, expected in test_data:
         check_consistency(
-            real_encoder.soft_real_encoder, real_encoder.hard_real_encoder, expected, input[0], input[1]
+            real_encoder.soft_real_encoder,
+            real_encoder.hard_real_encoder,
+            expected,
+            input[0],
+            input[1],
         )
 
 
@@ -50,7 +54,6 @@ def test_neuron():
         [0.0, [0.0, 0.0, 0.9], [0.5, 0.5, 0.0]],
         [1.0, [0.0, 1.0, 0.1], [1.0, 0.5, 1.0]],
         [0.0, [1.0, 0.0, 0.3], [0.0, 0.5, 0.0]],
-
         [0.3, [0.2, 0.8, 0.3], [0.5625, 0.1875, 0.5]],
         [0.1, [0.9, 0.42, 0.5], [0.05555556, 0.11904762, 0.1]],
         [0.4, [0.2, 0.8, 0.7], [0.625, 0.25, 0.2857143]],
@@ -108,3 +111,85 @@ def hard(thresholds, input):
             jax.numpy.array(input),
         )
 
+
+def test_real_encoder():
+    def test_net(type, x):
+        return real_encoder.real_encoder_layer(type)(3)(x)
+
+    soft, hard, symbolic = neural_logic_net.net(test_net)
+    weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
+    hard_weights = harden.hard_weights(weights)
+    print(f'weights: {weights}')
+    print(f'hard_weights: {hard_weights}')
+
+    test_data = [
+        [
+            [1.0, 0.8],
+            [
+                [1.0, 1.0, 1.0],
+                [0.47898874, 0.4623352,  0.6924789]
+            ],
+        ],
+        [
+            [0.6, 0.0],
+            [
+                [
+                    0.9469013,
+                    0.320184,
+                    0.3194083,
+                ],
+                [
+                    0.58414006,
+                    0.7815013,
+                    0.04193211,
+                ],
+            ],
+        ],
+        [
+            [0.1, 0.9],
+            [
+                [
+                    0.05309868,
+                    0.679816,
+                    0.6805917,
+                ],
+                [
+                    0.41585994,
+                    0.2184987,
+                    0.9580679,
+                ],
+            ],
+        ],
+        [
+            [0.4, 0.6],
+            [
+                [
+                    0.05309868,
+                    0.320184,
+                    0.6805917,
+                ],
+                [
+                    0.58414006,
+                    0.2184987,
+                    0.04193211,
+                ],
+            ],
+        ],
+    ]
+    for input, expected in test_data:
+        # Check that the soft function performs as expected
+        soft_output = soft.apply(weights, jax.numpy.array(input))
+        soft_expected = jax.numpy.array(expected)
+        print(f'soft_output: {soft_output}\nsoft_expected: {soft_expected}')
+        assert jax.numpy.allclose(soft_output, soft_expected)
+
+        # Check that the hard function performs as expected
+        hard_expected = harden.harden(jax.numpy.array(expected))
+        hard_output = hard.apply(hard_weights, jax.numpy.array(input))
+        print(f'hard_output: {hard_output}\nhard_expected: {hard_expected}')
+        assert jax.numpy.allclose(hard_output, hard_expected)
+
+        # Check that the symbolic function performs as expected
+        symbolic_output = symbolic.apply(hard_weights, jax.numpy.array(input))
+        print(f'symbolic_output: {symbolic_output}\nhard_expected: {hard_expected}')
+        assert numpy.allclose(symbolic_output, hard_expected)