Merge pull request #53 from github/remove-old-symbolic-eval

New symbolic evaluator

Author: Z80coder

neurallogic/hard_and.py

@@ -1,71 +1,49 @@
-from functools import reduce
-from typing import Callable
+from typing import Any
 
+import numpy
 import jax
 from flax import linen as nn
+from typing import Callable
+
 
-from neurallogic import neural_logic_net
+from neurallogic import neural_logic_net, symbolic_generation
 
 
 def soft_and_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: x OR ! w
     """
     w = jax.numpy.clip(w, 0.0, 1.0)
     return jax.numpy.maximum(x, 1.0 - w)
 
-@jax.jit
-def hard_and_include(w: bool, x: bool) -> bool:
-    return x | ~w
 
-def symbolic_and_include(w, x):
-    expression = f"({x} or not({w}))"
-    # Check if w is of type bool
-    if isinstance(w, bool) and isinstance(x, bool):
-        # We know the value of w and x, so we can evaluate the expression
-        return eval(expression)
-    # We don't know the value of w or x, so we return the expression
-    return expression
+
+def hard_and_include(w, x):
+    return jax.numpy.logical_or(x, jax.numpy.logical_not(w))
+
+
 
 def soft_and_neuron(w, x):
     x = jax.vmap(soft_and_include, 0, 0)(w, x)
     return jax.numpy.min(x)
 
+
 def hard_and_neuron(w, x):
     x = jax.vmap(hard_and_include, 0, 0)(w, x)
     return jax.lax.reduce(x, True, jax.lax.bitwise_and, [0])
 
-def symbolic_and_neuron(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    y = [symbolic_and_include(wi, xi) for wi, xi in zip(w, x)]
-    expression = "(" + str(reduce(lambda a, b: f"{a} and {b}", y)) + ")"
-    if all(isinstance(yi, bool) for yi in y):
-        # We know the value of all yis, so we can evaluate the expression
-        return eval(expression)
-    return expression
 
 soft_and_layer = jax.vmap(soft_and_neuron, (0, None), 0)
 
 hard_and_layer = jax.vmap(hard_and_neuron, (0, None), 0)
 
-def symbolic_and_layer(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    return [symbolic_and_neuron(wi, x) for wi in w]
 
-# TODO: investigate better initialization
 def initialize_near_to_zero():
+    # TODO: investigate better initialization
     def init(key, shape, dtype):
         dtype = jax.dtypes.canonicalize_dtype(dtype)
         # Sample from standard normal distribution (zero mean, unit variance)
@@ -76,6 +54,7 @@ def init(key, shape, dtype):
         return x
     return init
 
+
 class SoftAndLayer(nn.Module):
     """
     A soft-bit AND layer than transforms its inputs along the last dimension.
@@ -91,10 +70,12 @@ 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_shape, self.dtype)
+        weights = self.param('weights', self.weights_init,
+                             weights_shape, self.dtype)
         x = jax.numpy.asarray(x, self.dtype)
         return soft_and_layer(weights, x)
 
+
 class HardAndLayer(nn.Module):
     """
     A hard-bit And layer that shadows the SoftAndLayer.
@@ -108,26 +89,22 @@ class HardAndLayer(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(
+            'weights', nn.initializers.constant(0.0), weights_shape)
         return hard_and_layer(weights, x)
 
-class SymbolicAndLayer(nn.Module):
-    """A symbolic And layer than transforms its inputs along the last dimension.
-    Attributes:
-        layer_size: The number of neurons in the layer.
-    """
-    layer_size: int
 
-    @nn.compact
+class SymbolicAndLayer:
+    def __init__(self, layer_size):
+        self.layer_size = layer_size
+        self.hard_and_layer = HardAndLayer(self.layer_size)
+
     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 = weights.tolist()
-        if not isinstance(x, list):
-            raise TypeError(f"Input {x} should be a list")
-        return symbolic_and_layer(weights, x)
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(self.hard_and_layer, x)
+        return symbolic_generation.symbolic_expression(jaxpr, 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=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))

neurallogic/hard_not.py

@@ -3,7 +3,7 @@
 import jax
 from flax import linen as nn
 
-from neurallogic import neural_logic_net
+from neurallogic import neural_logic_net, symbolic_generation
 
 
 def soft_not(w: float, x: float) -> float:
@@ -18,57 +18,25 @@ def soft_not(w: float, x: float) -> float:
     return 1.0 - w + x * (2.0 * w - 1.0)
 
 
-@jax.jit
 def hard_not(w: bool, x: bool) -> bool:
-    return ~(x ^ w)
-
-
-def symbolic_not(w, x):
-    expression = f"(not({x} ^ {w}))"
-    # Check if w is of type bool
-    if isinstance(w, bool) and isinstance(x, bool):
-        # We know the value of w and x, so we can evaluate the expression
-        return eval(expression)
-    # We don't know the value of w or x, so we return the expression
-    return expression
+    # ~(x ^ w)
+    return jax.numpy.logical_not(jax.numpy.logical_xor(x, w))
 
 
 soft_not_neuron = jax.vmap(soft_not, 0, 0)
 
 hard_not_neuron = jax.vmap(hard_not, 0, 0)
 
 
-def symbolic_not_neuron(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    return [symbolic_not(wi, xi) for wi, xi in zip(w, x)]
-
 
 soft_not_layer = jax.vmap(soft_not_neuron, (0, None), 0)
 
 hard_not_layer = jax.vmap(hard_not_neuron, (0, None), 0)
 
 
-def symbolic_not_layer(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    return [symbolic_not_neuron(wi, x) for wi in w]
 
 
 class SoftNotLayer(nn.Module):
-    """
-    A soft-bit NOT 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.
-    """
     layer_size: int
     weights_init: Callable = nn.initializers.uniform(1.0)
     dtype: jax.numpy.dtype = jax.numpy.float32
@@ -83,13 +51,6 @@ def __call__(self, x):
 
 
 class HardNotLayer(nn.Module):
-    """
-    A hard-bit NOT layer that shadows the SoftNotLayer.
-    This is a convenience class to make it easier to switch between soft and hard logic.
-
-    Attributes:
-        layer_size: The number of neurons in the layer.
-    """
     layer_size: int
 
     @nn.compact
@@ -100,22 +61,14 @@ def __call__(self, x):
         return hard_not_layer(weights, x)
 
 
-class SymbolicNotLayer(nn.Module):
-    """A symbolic NOT layer than transforms its inputs along the last dimension.
-    Attributes:
-        layer_size: The number of neurons in the layer.
-    """
-    layer_size: int
+class SymbolicNotLayer:
+    def __init__(self, layer_size):
+        self.layer_size = layer_size
+        self.hard_not_layer = HardNotLayer(self.layer_size)
 
-    @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 = weights.tolist()
-        if not isinstance(x, list):
-            raise TypeError(f"Input {x} should be a list")
-        return symbolic_not_layer(weights, x)
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(self.hard_not_layer, x)
+        return symbolic_generation.symbolic_expression(jaxpr, x)
 
 
 not_layer = neural_logic_net.select(

neurallogic/hard_or.py

@@ -4,7 +4,7 @@
 import jax
 from flax import linen as nn
 
-from neurallogic import neural_logic_net
+from neurallogic import neural_logic_net, symbolic_generation
 
 
 def soft_or_include(w: float, x: float) -> float:
@@ -18,18 +18,10 @@ def soft_or_include(w: float, x: float) -> float:
     w = jax.numpy.clip(w, 0.0, 1.0)
     return 1.0 - jax.numpy.maximum(1.0 - x, 1.0 - w)
 
-@jax.jit
-def hard_or_include(w: bool, x: bool) -> bool:
-    return x & w
 
-def symbolic_or_include(w, x):
-    expression = f"({x} and {w})"
-    # Check if w is of type bool
-    if isinstance(w, bool) and isinstance(x, bool):
-        # We know the value of w and x, so we can evaluate the expression
-        return eval(expression)
-    # We don't know the value of w or x, so we return the expression
-    return expression
+def hard_or_include(w, x):
+    return jax.numpy.logical_and(x, w)
+
 
 def soft_or_neuron(w, x):
     x = jax.vmap(soft_or_include, 0, 0)(w, x)
@@ -39,31 +31,11 @@ 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])
 
-def symbolic_or_neuron(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    y = [symbolic_or_include(wi, xi) for wi, xi in zip(w, x)]
-    expression = "(" + str(reduce(lambda a, b: f"{a} or {b}", y)) + ")"
-    if all(isinstance(yi, bool) for yi in y):
-        # We know the value of all yis, so we can evaluate the expression
-        return eval(expression)
-    return expression
 
 soft_or_layer = jax.vmap(soft_or_neuron, (0, None), 0)
 
 hard_or_layer = jax.vmap(hard_or_neuron, (0, None), 0)
 
-def symbolic_or_layer(w, x):
-    # TODO: ensure that this implementation has the same generality over tensors as vmap
-    if not isinstance(w, list):
-        raise TypeError(f"Input {x} should be a list")
-    if not isinstance(x, list):
-        raise TypeError(f"Input {x} should be a list")
-    return [symbolic_or_neuron(wi, x) for wi in w]
-
 # TODO: investigate better initialization
 def initialize_near_to_one():
     def init(key, shape, dtype):
@@ -77,13 +49,6 @@ def init(key, shape, dtype):
     return init
 
 class SoftOrLayer(nn.Module):
-    """
-    A soft-bit Or 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.
-    """
     layer_size: int
     weights_init: Callable = initialize_near_to_one()
     dtype: jax.numpy.dtype = jax.numpy.float32
@@ -96,13 +61,6 @@ def __call__(self, x):
         return soft_or_layer(weights, x)
 
 class HardOrLayer(nn.Module):
-    """
-    A hard-bit Or layer that shadows the SoftAndLayer.
-    This is a convenience class to make it easier to switch between soft and hard logic.
-
-    Attributes:
-        layer_size: The number of neurons in the layer.
-    """
     layer_size: int
 
     @nn.compact
@@ -111,21 +69,14 @@ def __call__(self, x):
         weights = self.param('weights', nn.initializers.constant(0.0), weights_shape)
         return hard_or_layer(weights, x)
 
-class SymbolicOrLayer(nn.Module):
-    """A symbolic Or layer than transforms its inputs along the last dimension.
-    Attributes:
-        layer_size: The number of neurons in the layer.
-    """
-    layer_size: int
+class SymbolicOrLayer:
+    def __init__(self, layer_size):
+        self.layer_size = layer_size
+        self.hard_or_layer = HardOrLayer(self.layer_size)
 
-    @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 = weights.tolist()
-        if not isinstance(x, list):
-            raise TypeError(f"Input {x} should be a list")
-        return symbolic_or_layer(weights, 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),