get more of test_hard_not working

Author: Z80coder

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,24 @@ 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
+    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 +50,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 +60,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/symbolic_primitives.py

@@ -155,6 +155,9 @@ def binary_infix_operator(operator: str, a: numpy.ndarray, b: float, bracket: bo
 def binary_infix_operator(operator: str, a: str, b: float, bracket: bool = False):
     return binary_infix_operator(operator, a, str(b), bracket)
 
+@dispatch
+def binary_infix_operator(operator: str, a: numpy.ndarray, b: jax.numpy.ndarray, bracket: bool = False):
+    return binary_infix_operator(operator, a, numpy.array(b), bracket)
 
 
 def all_concrete_values(data):

tests/test_hard_and.py

@@ -121,7 +121,7 @@ def test_net(type, x):
         return hard_and.and_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],
@@ -141,16 +141,16 @@ def test_net(type, x):
     # 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: 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)
+    weights = state.params
+    hard_weights = harden.hard_weights(weights)
 
     for input, expected in zip(x, y):
         hard_input = harden.harden(jax.numpy.array(input))
@@ -171,11 +171,11 @@ def test_net(type, x):
 
     # Compute soft result
     soft_input = jax.numpy.array([0.6, 0.45])
-    soft_weights = soft.init(random.PRNGKey(0), soft_input)
-    soft_result = soft.apply(soft_weights, numpy.array(soft_input))
+    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(soft_weights)
+    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