add experimental xor neuron

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_not.py

@@ -48,11 +48,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 +68,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)
+        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,96 @@
+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)
+    # Compute the most sensitive bit
+    margins = jax.vmap(lambda x: jax.numpy.abs(0.5 - x))(x)
+    sensitive_bit_index = jax.numpy.argmin(margins)
+    sensitive_bit = jax.numpy.take(x, sensitive_bit_index)
+    # Compute the logical xor of the bits
+    hard_x = jax.vmap(lambda x: jax.numpy.where(x > 0.5, True, False))(x)
+    logical_xor = jax.lax.reduce(hard_x, False, jax.numpy.logical_xor, (0,))
+    # Compute the representative bit
+    hard_sensitive_bit = jax.numpy.where(sensitive_bit > 0.5, True, False)
+    representative_bit = jax.numpy.where(logical_xor == hard_sensitive_bit,
+                                         sensitive_bit,
+                                         1.0 - sensitive_bit
+                                         )
+    return representative_bit
+
+
+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)  # TODO: investigate better initialization
+    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/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

neurallogic/symbolic_generation.py

@@ -41,6 +41,7 @@ def symbolic_bind(prim, *args, **params):
         'not': symbolic_primitives.symbolic_not,
         'reduce_and': symbolic_primitives.symbolic_reduce_and,
         'reduce_or': symbolic_primitives.symbolic_reduce_or,
+        'reduce_xor': symbolic_primitives.symbolic_reduce_xor,
         'reduce_sum': symbolic_primitives.symbolic_reduce_sum,
         'select_n': symbolic_primitives.symbolic_select_n,
     }[prim.name](*args, **params)

neurallogic/symbolic_primitives.py

@@ -262,6 +262,23 @@ def symbolic_reduce_or(*args, **kwargs):
         )
 
 
+def symbolic_reduce_xor(*args, **kwargs):
+    if all_concrete_values([*args]):
+        return lax_reference.reduce(
+            *args,
+            init_value=False,
+            computation=numpy.logical_xor,
+            dimensions=kwargs['axes'],
+        )
+    else:
+        return symbolic_reduce(
+            *args,
+            init_value='False',
+            computation=symbolic_xor,
+            dimensions=kwargs['axes'],
+        )
+
+
 def symbolic_reduce_sum(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.reduce(

tests/test_hard_or.py

@@ -1,5 +1,4 @@
 import jax
-import jax.numpy as jnp
 import numpy
 import optax
 from flax import linen as nn
@@ -134,7 +133,7 @@ def test_net(type, x):
     input = jax.numpy.array(x)
     output = jax.numpy.array(y)
 
-    # Train the and layer
+    # Train the or layer
     tx = optax.sgd(0.1)
     state = train_state.TrainState.create(apply_fn=jax.vmap(
         soft.apply, in_axes=(None, 0)), params=weights, tx=tx)