Merge pull request #54 from github/mnist-exp

Real encoder

Author: Z80coder

neurallogic/hard_and.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import numpy
 import jax
 from flax import linen as nn
 from typing import Callable
@@ -70,7 +69,7 @@ class SoftAndLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', self.weights_init,
+        weights = self.param('bit_weights', self.weights_init,
                              weights_shape, self.dtype)
         x = jax.numpy.asarray(x, self.dtype)
         return soft_and_layer(weights, x)
@@ -90,7 +89,7 @@ class HardAndLayer(nn.Module):
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'weights', nn.initializers.constant(0.0), weights_shape)
+            'bit_weights', nn.initializers.constant(0.0), weights_shape)
         return hard_and_layer(weights, x)
 
 

neurallogic/hard_not.py

@@ -28,14 +28,11 @@ 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, None), 0)
 
 hard_not_layer = jax.vmap(hard_not_neuron, (0, None), 0)
 
 
-
-
 class SoftNotLayer(nn.Module):
     layer_size: int
     weights_init: Callable = nn.initializers.uniform(1.0)
@@ -44,8 +41,7 @@ class SoftNotLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', self.weights_init,
-                             weights_shape, self.dtype)
+        weights = self.param("bit_weights", self.weights_init, weights_shape, self.dtype)
         x = jax.numpy.asarray(x, self.dtype)
         return soft_not_layer(weights, x)
 
@@ -56,8 +52,7 @@ class HardNotLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param(
-            'weights', nn.initializers.constant(0.0), weights_shape)
+        weights = self.param("bit_weights", nn.initializers.constant(0.0), weights_shape)
         return hard_not_layer(weights, x)
 
 
@@ -73,7 +68,12 @@ def __call__(self, x):
 
 not_layer = neural_logic_net.select(
     lambda layer_size, weights_init=nn.initializers.uniform(
-        1.0), dtype=jax.numpy.float32: SoftNotLayer(layer_size, weights_init, dtype),
+        1.0
+    ), dtype=jax.numpy.float32: SoftNotLayer(layer_size, weights_init, dtype),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: HardNotLayer(layer_size),
     lambda layer_size, weights_init=nn.initializers.uniform(
-        1.0), dtype=jax.numpy.float32: HardNotLayer(layer_size),
-    lambda layer_size, weights_init=nn.initializers.uniform(1.0), dtype=jax.numpy.float32: SymbolicNotLayer(layer_size))
+        1.0
+    ), dtype=jax.numpy.float32: SymbolicNotLayer(layer_size),
+)

neurallogic/hard_or.py

@@ -56,7 +56,7 @@ class SoftOrLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', self.weights_init, weights_shape, self.dtype)
+        weights = self.param('bit_weights', self.weights_init, weights_shape, self.dtype)
         x = jax.numpy.asarray(x, self.dtype)
         return soft_or_layer(weights, x)
 
@@ -66,7 +66,7 @@ class HardOrLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', nn.initializers.constant(0.0), weights_shape)
+        weights = self.param('bit_weights', nn.initializers.constant(0.0), weights_shape)
         return hard_or_layer(weights, x)
 
 class SymbolicOrLayer:

neurallogic/harden.py

@@ -2,47 +2,74 @@
 import jax
 import numpy
 from plum import dispatch
-from neurallogic import symbolic_primitives
+
+from neurallogic import map_at_elements
 
 
 def harden_float(x: float) -> bool:
     return x > 0.5
 
+
 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)
+    return map_at_elements.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)
+    # Only harden parameters that explicitly represent bits
+    def conditional_harden(k, v):
+        if k.startswith("bit_"):
+            return map_at_elements.map_at_elements(v, harden)
+        elif isinstance(v, dict) or isinstance(v, flax.core.FrozenDict) or isinstance(v, list):
+            return harden(v)
+        return v
+
+    return {k: conditional_harden(k, v) for k, v in x.items()}
+
 
 @dispatch
 def harden(x: flax.core.FrozenDict):
-    return flax.core.FrozenDict(symbolic_primitives.map_at_elements(x.unfreeze(), harden_float))
+    return harden(x.unfreeze())
+
 
 @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/map_at_elements.py

@@ -0,0 +1,58 @@
+import typing
+
+import jax
+import numpy
+from plum import dispatch
+
+
+@dispatch
+def map_at_elements(x: str, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: bool, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: numpy.bool_, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: float, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: numpy.float32, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: list, func: typing.Callable):
+    return [map_at_elements(item, func) for item in x]
+
+
+@dispatch
+def map_at_elements(x: numpy.ndarray, func: typing.Callable):
+    return numpy.array([map_at_elements(item, func) for item in x], dtype=object)
+
+
+@dispatch
+def map_at_elements(x: jax.numpy.ndarray, func: typing.Callable):
+    if x.ndim == 0:
+        return func(x.item())
+    return jax.numpy.array([map_at_elements(item, func) for item in x])
+
+
+@dispatch
+def map_at_elements(x: dict, func: typing.Callable):
+    return {k: map_at_elements(v, func) for k, v in x.items()}
+
+
+@dispatch
+def map_at_elements(x: tuple, func: typing.Callable):
+    return tuple(map_at_elements(list(x), func))
+

neurallogic/real_encoder.py

@@ -0,0 +1,89 @@
+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:
+    eps = 0.0000001
+    # x should be in [0, 1]
+    t = jax.numpy.clip(t, 0.0, 1.0)
+    return jax.numpy.where(
+        jax.numpy.isclose(t, x),
+        0.5,
+        # t != x
+        jax.numpy.where(
+            x < t,
+            (1.0 / (2.0 * t + eps)) * x,
+            # x > t
+            (1.0 / (2.0 * (1.0 - t) + eps)) * (x + 1.0 - 2.0 * t)
+        )
+    )
+
+
+def hard_real_encoder(t: float, x: float) -> bool:
+    # t and x must be floats
+    return jax.numpy.where(soft_real_encoder(t, x) > 0.5, True, False)
+
+
+soft_real_encoder_neuron = jax.vmap(soft_real_encoder, in_axes=(0, None))
+
+hard_real_encoder_neuron = jax.vmap(hard_real_encoder, in_axes=(0, None))
+
+soft_real_encoder_layer = jax.vmap(soft_real_encoder_neuron, (0, 0), 0)
+
+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),
+)