Merge pull request #55 from github/mnist-2

Additional layer types

Author: Z80coder

neurallogic/hard_and.py

@@ -89,7 +89,7 @@ class HardAndLayer(nn.Module):
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'bit_weights', nn.initializers.constant(0.0), weights_shape)
+            'bit_weights', nn.initializers.constant(True), weights_shape)
         return hard_and_layer(weights, x)
 
 
@@ -105,5 +105,5 @@ def __call__(self, x):
 
 and_layer = neural_logic_net.select(
     lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: SoftAndLayer(layer_size, weights_init, dtype),
-    lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: HardAndLayer(layer_size),
-    lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: SymbolicAndLayer(layer_size))
+    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: HardAndLayer(layer_size),
+    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: SymbolicAndLayer(layer_size))

neurallogic/hard_majority.py

@@ -0,0 +1,60 @@
+import jax
+from flax import linen as nn
+
+from neurallogic import neural_logic_net, symbolic_generation
+
+
+def majority_index(input_size: int) -> int:
+    return (input_size - 1) // 2
+
+
+def soft_majority(x: jax.numpy.array) -> float:
+    index = majority_index(x.shape[-1])
+    sorted_x = jax.numpy.sort(x, axis=-1)
+    return jax.numpy.take(sorted_x, index, axis=-1)
+
+
+def hard_majority(x: jax.numpy.array) -> bool:
+    threshold = x.shape[-1] - majority_index(x.shape[-1])
+    return jax.numpy.sum(x, axis=-1) >= threshold
+
+
+soft_majority_layer = jax.vmap(soft_majority, in_axes=0)
+
+hard_majority_layer = jax.vmap(hard_majority, in_axes=0)
+
+
+class SoftMajorityLayer(nn.Module):
+    """
+    A soft-bit MAJORITY layer than transforms its inputs along the last dimension.
+
+    Attributes:
+        layer_size: The number of neurons in the layer.
+        weights_init: The initializer function for the weight matrix.
+    """
+    @nn.compact
+    def __call__(self, x):
+        return soft_majority_layer(x)
+
+
+class HardMajorityLayer(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        return hard_majority_layer(x)
+
+
+class SymbolicMajorityLayer:
+    def __init__(self):
+        self.hard_majority_layer = HardMajorityLayer()
+
+    def __call__(self, x):
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
+            self.hard_majority_layer, x)
+        return symbolic_generation.symbolic_expression(jaxpr, x)
+
+
+majority_layer = neural_logic_net.select(
+    lambda: SoftMajorityLayer(),
+    lambda: HardMajorityLayer(),
+    lambda: SymbolicMajorityLayer()
+)

neurallogic/hard_not.py

@@ -19,7 +19,6 @@ def soft_not(w: float, x: float) -> float:
 
 
 def hard_not(w: bool, x: bool) -> bool:
-    # ~(x ^ w)
     return jax.numpy.logical_not(jax.numpy.logical_xor(x, w))
 
 
@@ -48,11 +47,12 @@ def __call__(self, x):
 
 class HardNotLayer(nn.Module):
     layer_size: int
+    weights_init: Callable = nn.initializers.constant(True)
 
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param("bit_weights", nn.initializers.constant(0.0), weights_shape)
+        weights = self.param("bit_weights", self.weights_init, weights_shape)
         return hard_not_layer(weights, x)
 
 
@@ -67,13 +67,7 @@ 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),
-    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),
+    lambda layer_size, weights_init=nn.initializers.uniform(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: SymbolicNotLayer(layer_size),
 )

neurallogic/hard_or.py

@@ -1,4 +1,3 @@
-from functools import reduce
 from typing import Callable
 
 import jax
@@ -12,7 +11,7 @@ def soft_or_include(w: float, x: float) -> float:
     w > 0.5 implies the and operation is active, else inactive
 
     Assumes x is in [0, 1]
-    
+
     Corresponding hard logic: b AND w
     """
     w = jax.numpy.clip(w, 0.0, 1.0)
@@ -27,6 +26,7 @@ def soft_or_neuron(w, x):
     x = jax.vmap(soft_or_include, 0, 0)(w, x)
     return jax.numpy.max(x)
 
+
 def hard_or_neuron(w, x):
     x = jax.vmap(hard_or_include, 0, 0)(w, x)
     return jax.lax.reduce(x, False, jax.lax.bitwise_or, [0])
@@ -37,6 +37,8 @@ def hard_or_neuron(w, x):
 hard_or_layer = jax.vmap(hard_or_neuron, (0, None), 0)
 
 # TODO: investigate better initialization
+
+
 def initialize_near_to_one():
     def init(key, shape, dtype):
         dtype = jax.dtypes.canonicalize_dtype(dtype)
@@ -48,6 +50,7 @@ def init(key, shape, dtype):
         return x
     return init
 
+
 class SoftOrLayer(nn.Module):
     layer_size: int
     weights_init: Callable = initialize_near_to_one()
@@ -56,29 +59,37 @@ class SoftOrLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('bit_weights', self.weights_init, weights_shape, self.dtype)
-        x = jax.numpy.asarray(x, 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)
 
+
 class HardOrLayer(nn.Module):
     layer_size: int
 
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('bit_weights', nn.initializers.constant(0.0), weights_shape)
+        weights = self.param(
+            'bit_weights', nn.initializers.constant(True), weights_shape)
         return hard_or_layer(weights, x)
 
+
 class SymbolicOrLayer:
     def __init__(self, layer_size):
         self.layer_size = layer_size
         self.hard_or_layer = HardOrLayer(self.layer_size)
 
     def __call__(self, x):
-        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(self.hard_or_layer, x)
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
+            self.hard_or_layer, x)
         return symbolic_generation.symbolic_expression(jaxpr, x)
 
+
 or_layer = neural_logic_net.select(
-    lambda layer_size, weights_init=initialize_near_to_one(), dtype=jax.numpy.float32: SoftOrLayer(layer_size, weights_init, dtype),
-    lambda layer_size, weights_init=initialize_near_to_one(), dtype=jax.numpy.float32: HardOrLayer(layer_size),
-    lambda layer_size, weights_init=initialize_near_to_one(), dtype=jax.numpy.float32: SymbolicOrLayer(layer_size))
+    lambda layer_size, weights_init=initialize_near_to_one(
+    ), dtype=jax.numpy.float32: SoftOrLayer(layer_size, weights_init, dtype),
+    lambda layer_size, weights_init=nn.initializers.constant(
+        True), dtype=jax.numpy.float32: HardOrLayer(layer_size),
+    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: SymbolicOrLayer(layer_size))

neurallogic/hard_xor.py

@@ -0,0 +1,87 @@
+from typing import Callable
+
+import jax
+from flax import linen as nn
+
+from neurallogic import neural_logic_net, symbolic_generation
+
+
+def soft_xor_include(w: float, x: float) -> float:
+    """
+    w > 0.5 implies the and operation is active, else inactive
+
+    Assumes x is in [0, 1]
+
+    Corresponding hard logic: b AND w
+    """
+    w = jax.numpy.clip(w, 0.0, 1.0)
+    return 1.0 - jax.numpy.maximum(1.0 - x, 1.0 - w)
+
+
+def hard_xor_include(w, x):
+    return jax.numpy.logical_and(x, w)
+
+
+def soft_xor_neuron(w, x):
+    # Conditionally include input bits, according to weights
+    x = jax.vmap(soft_xor_include, 0, 0)(w, x)
+
+    def xor(x, y):
+        return jax.numpy.minimum(jax.numpy.maximum(x, y), 1.0 - jax.numpy.minimum(x, y))
+    x = jax.lax.reduce(x, jax.numpy.array(0, dtype=x.dtype), xor, (0,))
+    return x
+
+
+def hard_xor_neuron(w, x):
+    x = jax.vmap(hard_xor_include, 0, 0)(w, x)
+    return jax.lax.reduce(x, False, jax.lax.bitwise_xor, [0])
+
+
+soft_xor_layer = jax.vmap(soft_xor_neuron, (0, None), 0)
+
+
+hard_xor_layer = jax.vmap(hard_xor_neuron, (0, None), 0)
+
+
+class SoftXorLayer(nn.Module):
+    layer_size: int
+    weights_init: Callable = nn.initializers.uniform(1.0)
+    dtype: jax.numpy.dtype = jax.numpy.float32
+
+    @nn.compact
+    def __call__(self, x):
+        weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
+        weights = self.param(
+            'bit_weights', self.weights_init, weights_shape, self.dtype)
+        x = jax.numpy.asarray(x, self.dtype) 
+        return soft_xor_layer(weights, x)
+
+
+class HardXorLayer(nn.Module):
+    layer_size: int
+
+    @nn.compact
+    def __call__(self, x):
+        weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
+        weights = self.param(
+            'bit_weights', nn.initializers.constant(True), weights_shape)
+        return hard_xor_layer(weights, x)
+
+
+class SymbolicXorLayer:
+    def __init__(self, layer_size):
+        self.layer_size = layer_size
+        self.hard_xor_layer = HardXorLayer(self.layer_size)
+
+    def __call__(self, x):
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
+            self.hard_xor_layer, x)
+        return symbolic_generation.symbolic_expression(jaxpr, x)
+
+
+xor_layer = neural_logic_net.select(
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0), dtype=jax.numpy.float32: SoftXorLayer(layer_size, weights_init, dtype),
+    lambda layer_size, weights_init=nn.initializers.constant(
+        True), dtype=jax.numpy.float32: HardXorLayer(layer_size),
+    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: SymbolicXorLayer(layer_size))

neurallogic/harden.py

@@ -20,18 +20,27 @@ def harden(x: float):
     return harden_float(x)
 
 
+@dispatch
+def harden(x: bool):
+    return x
+
+
 @dispatch
 def harden(x: list):
     return map_at_elements.map_at_elements(x, harden_float)
 
 
 @dispatch
 def harden(x: numpy.ndarray):
+    if x.ndim == 0:
+        return harden(x.item())
     return harden_array(x)
 
 
 @dispatch
 def harden(x: jax.numpy.ndarray):
+    if x.ndim == 0:
+        return harden(x.item())
     return harden_array(x)
 
 
@@ -53,13 +62,6 @@ def harden(x: flax.core.FrozenDict):
     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 {

neurallogic/harden_layer.py

@@ -18,6 +18,8 @@ def straight_through_harden_element(x):
 def hard_harden_layer(x):
     return x
 
+#TODO: can we harden arbitrary tensors?
+#TODO: is this correct?
 def symbolic_harden_layer(x):
     return x
 

neurallogic/map_at_elements.py

@@ -30,6 +30,11 @@ def map_at_elements(x: numpy.float32, func: typing.Callable):
     return func(x)
 
 
+@dispatch
+def map_at_elements(x: numpy.int32, 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]
@@ -55,4 +60,3 @@ def map_at_elements(x: dict, func: typing.Callable):
 @dispatch
 def map_at_elements(x: tuple, func: typing.Callable):
     return tuple(map_at_elements(list(x), func))
-

neurallogic/real_encoder.py

@@ -5,6 +5,8 @@
 
 from neurallogic import neural_logic_net, symbolic_generation
 
+# TODO: perhaps this can be simplified with a simple multiplication?
+# TODO: implement a soft_real_decoder that can perhaps replace the port count approach
 
 def soft_real_encoder(t: float, x: float) -> float:
     eps = 0.0000001