get hard_or test passing

Author: Z80coder

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),

tests/test_hard_and.py

@@ -4,7 +4,6 @@
 from flax.training import train_state
 from jax import random
 import numpy
-import typing
 
 from neurallogic import hard_and, harden, neural_logic_net, symbolic_generation
 from tests import utils

tests/test_hard_or.py

@@ -4,8 +4,10 @@
 from flax import linen as nn
 from flax.training import train_state
 from jax import random
+import numpy
 
-from neurallogic import hard_or, harden, neural_logic_net, primitives
+from neurallogic import hard_or, harden, neural_logic_net, symbolic_generation
+from tests import utils
 
 
 def test_include():
@@ -20,11 +22,8 @@ def test_include():
         [[-0.1, 1.0], 0.0]
     ]
     for input, expected in test_data:
-        assert hard_or.soft_or_include(*input) == expected
-        assert hard_or.hard_or_include(
-            *harden.harden(input)) == harden.harden(expected)
-        symbolic_output = hard_or.symbolic_or_include(*harden.harden(input))
-        assert symbolic_output == harden.harden(expected)
+        utils.check_consistency(hard_or.soft_or_include, hard_or.hard_or_include,
+                           expected, input[0], input[1])
 
 
 def test_neuron():
@@ -37,14 +36,14 @@ def test_neuron():
         [[0.0, 1.0], [1.0, 1.0], 1.0]
     ]
     for input, weights, expected in test_data:
-        input = jnp.array(input)
-        weights = jnp.array(weights)
-        assert jnp.allclose(hard_or.soft_or_neuron(weights, input), expected)
-        assert jnp.allclose(hard_or.hard_or_neuron(harden.harden(
-            weights), harden.harden(input)), harden.harden(expected))
-        symbolic_output = hard_or.symbolic_or_neuron(
-            harden.harden(weights.tolist()), harden.harden(input.tolist()))
-        assert jnp.array_equal(symbolic_output, harden.harden(expected))
+        def soft(weights, input):
+            return hard_or.soft_or_neuron(weights, input)
+
+        def hard(weights, input):
+            return hard_or.hard_or_neuron(weights, input)
+
+        utils.check_consistency(soft, hard, expected,
+                          jax.numpy.array(weights), jax.numpy.array(input))
 
 
 def test_layer():
@@ -59,27 +58,27 @@ def test_layer():
             1.0, 0.0], [0.0, 0.0]], [0.0, 0.0, 0.0, 0.0]]
     ]
     for input, weights, expected in test_data:
-        input = jnp.array(input)
-        weights = jnp.array(weights)
-        expected = jnp.array(expected)
-        assert jnp.allclose(hard_or.soft_or_layer(weights, input), expected)
-        assert jnp.allclose(hard_or.hard_or_layer(harden.harden(
-            weights), harden.harden(input)), harden.harden(expected))
-        symbolic_output = hard_or.symbolic_or_layer(
-            harden.harden(weights.tolist()), harden.harden(input.tolist()))
-        assert jnp.array_equal(symbolic_output, harden.harden(expected))
+        def soft(weights, input):
+            return hard_or.soft_or_layer(weights, input)
+
+        def hard(weights, input):
+            return hard_or.hard_or_layer(weights, input)
+
+
+        utils.check_consistency(soft, hard, jax.numpy.array(expected),
+                          jax.numpy.array(weights), jax.numpy.array(input))
 
 
 def test_or():
     def test_net(type, x):
         x = hard_or.or_layer(type)(4, nn.initializers.uniform(1.0))(x)
-        x = primitives.nl_ravel(type)(x)
+        x = x.ravel()
         return x
 
     soft, hard, symbolic = neural_logic_net.net(test_net)
-    soft_weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
-    hard_weights = harden.hard_weights(soft_weights)
-    symbolic_weights = harden.symbolic_weights(soft_weights)
+    weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
+    hard_weights = harden.hard_weights(weights)
+
     test_data = [
         [
             [1.0, 1.0],
@@ -99,24 +98,28 @@ def test_net(type, x):
         ]
     ]
     for input, expected in test_data:
-        soft_input = jnp.array(input)
-        soft_expected = jnp.array(expected)
-        soft_result = soft.apply(soft_weights, soft_input)
-        assert jnp.allclose(soft_result, soft_expected)
-        hard_input = harden.harden(soft_input)
-        hard_expected = harden.harden(soft_expected)
-        hard_result = hard.apply(hard_weights, hard_input)
-        assert jnp.allclose(hard_result, hard_expected)
-        symbolic_result = symbolic.apply(symbolic_weights, hard_input.tolist())
-        assert jnp.array_equal(symbolic_result, hard_expected)
+        # Check that the soft function performs as expected
+        assert jax.numpy.allclose(soft.apply(
+            weights, jax.numpy.array(input)), jax.numpy.array(expected))
+
+        # Check that the hard function performs as expected
+        hard_input = harden.harden(jax.numpy.array(input))
+        hard_expected = harden.harden(jax.numpy.array(expected))
+        hard_output = hard.apply(hard_weights, hard_input)
+        assert jax.numpy.allclose(hard_output, hard_expected)
+
+        # Check that the symbolic function performs as expected
+        symbolic_output = symbolic.apply(hard_weights, hard_input)
+        assert numpy.allclose(symbolic_output, hard_expected)
 
 
 def test_train_or():
     def test_net(type, x):
         return hard_or.or_layer(type)(4, nn.initializers.uniform(1.0))(x)
 
     soft, hard, symbolic = neural_logic_net.net(test_net)
-    soft_weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
+    weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
+
     x = [
         [1.0, 1.0],
         [1.0, 0.0],
@@ -129,30 +132,31 @@ def test_net(type, x):
         [1.0, 0.0, 0.0, 1.0],
         [0.0, 0.0, 0.0, 0.0]
     ]
-    input = jnp.array(x)
-    output = jnp.array(y)
+    input = jax.numpy.array(x)
+    output = jax.numpy.array(y)
 
     # Train the and layer
     tx = optax.sgd(0.1)
     state = train_state.TrainState.create(apply_fn=jax.vmap(
-        soft.apply, in_axes=(None, 0)), params=soft_weights, tx=tx)
+        soft.apply, in_axes=(None, 0)), params=weights, tx=tx)
     grad_fn = jax.jit(jax.value_and_grad(lambda params, x,
-                      y: jnp.mean((state.apply_fn(params, x) - y) ** 2)))
+                      y: jax.numpy.mean((state.apply_fn(params, x) - y) ** 2)))
     for epoch in range(1, 100):
         loss, grads = grad_fn(state.params, input, output)
         state = state.apply_gradients(grads=grads)
 
     # Test that the and layer (both soft and hard variants) correctly predicts y
-    soft_weights = state.params
-    hard_weights = harden.hard_weights(soft_weights)
-    symbolic_weights = harden.symbolic_weights(soft_weights)
+    weights = state.params
+    hard_weights = harden.hard_weights(weights)
+
     for input, expected in zip(x, y):
-        hard_input = harden.harden_array(harden.harden(jnp.array(input)))
-        hard_expected = harden.harden_array(harden.harden(jnp.array(expected)))
+        hard_input = harden.harden(jax.numpy.array(input))
+        hard_expected = harden.harden(jax.numpy.array(expected))
         hard_result = hard.apply(hard_weights, hard_input)
-        assert jnp.allclose(hard_result, hard_expected)
-        symbolic_result = symbolic.apply(symbolic_weights, hard_input.tolist())
-        assert jnp.array_equal(symbolic_result, hard_expected)
+        assert jax.numpy.allclose(hard_result, hard_expected)
+        symbolic_output = symbolic.apply(hard_weights, hard_input)
+        assert jax.numpy.array_equal(symbolic_output, hard_expected)
+
 
 
 def test_symbolic_or():
@@ -162,9 +166,46 @@ def test_net(type, x):
         return x
 
     soft, hard, symbolic = neural_logic_net.net(test_net)
-    soft_weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
-    symbolic_weights = harden.symbolic_weights(soft_weights)
+
+    # Compute soft result
+    soft_input = jax.numpy.array([0.6, 0.45])
+    weights = soft.init(random.PRNGKey(0), soft_input)
+    soft_result = soft.apply(weights, numpy.array(soft_input))
+    
+    # Compute hard result
+    hard_weights = harden.hard_weights(weights)
+    hard_input = harden.harden(soft_input)
+    hard_result = hard.apply(hard_weights, numpy.array(hard_input))
+    # Check that the hard result is the same as the soft result
+    assert numpy.array_equal(harden.harden(soft_result), hard_result)
+
+    # Compute symbolic result with non-symbolic inputs
+    symbolic_output = symbolic.apply(hard_weights, hard_input)
+    # Check that the symbolic result is the same as the hard result
+    assert numpy.array_equal(symbolic_output, hard_result)
+
+    # Compute symbolic result with symbolic inputs and symbolic weights, but where the symbols can be evaluated
+    symbolic_input = ['True', 'False']
+    symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+    symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
+    symbolic_output = symbolic_generation.eval_symbolic_expression(symbolic_output)
+    # Check that the symbolic result is the same as the hard result
+    assert numpy.array_equal(symbolic_output, hard_result)
+
+    # Compute symbolic result with symbolic inputs and non-symbolic weights
     symbolic_input = ['x1', 'x2']
-    symbolic_result = symbolic.apply(symbolic_weights, symbolic_input)
-    assert (symbolic_result == ['((((x1 and False) or (x2 and False)) and True) or (((x1 and False) or (x2 and False)) and True) or (((x1 and False) or (x2 and True)) and True) or (((x1 and True) or (x2 and True)) and True))', '((((x1 and False) or (x2 and False)) and False) or (((x1 and False) or (x2 and False)) and True) or (((x1 and False) or (x2 and True)) and False) or (((x1 and True) or (x2 and True)) and False))',
-            '((((x1 and False) or (x2 and False)) and False) or (((x1 and False) or (x2 and False)) and False) or (((x1 and False) or (x2 and True)) and True) or (((x1 and True) or (x2 and True)) and True))', '((((x1 and False) or (x2 and False)) and False) or (((x1 and False) or (x2 and False)) and True) or (((x1 and False) or (x2 and True)) and True) or (((x1 and True) or (x2 and True)) and True))'])
+    symbolic_output = symbolic.apply(hard_weights, symbolic_input)
+    # Check the form of the symbolic expression
+    assert numpy.array_equal(symbolic_output, ['((((False or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and True) or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and True) or (((False or (x1 != 0) and False) or (x2 != 0) and True) != 0) and True) or (((False or (x1 != 0) and True) or (x2 != 0) and True) != 0) and True)',
+ '((((False or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and False) or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and True) or (((False or (x1 != 0) and False) or (x2 != 0) and True) != 0) and False) or (((False or (x1 != 0) and True) or (x2 != 0) and True) != 0) and False)',
+ '((((False or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and False) or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and False) or (((False or (x1 != 0) and False) or (x2 != 0) and True) != 0) and True) or (((False or (x1 != 0) and True) or (x2 != 0) and True) != 0) and True)',
+ '((((False or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and False) or (((False or (x1 != 0) and False) or (x2 != 0) and False) != 0) and True) or (((False or (x1 != 0) and False) or (x2 != 0) and True) != 0) and True) or (((False or (x1 != 0) and True) or (x2 != 0) and True) != 0) and True)'])
+
+    # Compute symbolic result with symbolic inputs and symbolic weights
+    symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
+    # Check the form of the symbolic expression
+    assert numpy.array_equal(symbolic_output, ['((((False or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (True != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0))',
+ '((((False or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (False != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (True != 0)) != 0) and (False != 0)) or (((False or (x1 != 0) and (True != 0)) or (x2 != 0) and (True != 0)) != 0) and (False != 0))',
+ '((((False or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (False != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (False != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (True != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0))',
+ '((((False or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (False != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (False != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (False != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0)) or (((False or (x1 != 0) and (True != 0)) or (x2 != 0) and (True != 0)) != 0) and (True != 0))'])
+