symbolic version of majority, plus additonal tests

Author: Z80coder

neurallogic/hard_majority.py

@@ -1,6 +1,7 @@
 import jax
+from flax import linen as nn
 
-from neurallogic import neural_logic_net
+from neurallogic import neural_logic_net, symbolic_generation
 
 
 def majority_index(input_size: int) -> int:
@@ -23,9 +24,40 @@ def hard_majority(x: jax.numpy.array) -> bool:
 hard_majority_layer = jax.vmap(hard_majority, in_axes=0)
 
 
-def symbolic_majority_layer(x):
-    return hard_majority_layer(x)
+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.
+    """
+    dtype: jax.numpy.dtype = jax.numpy.float32
+
+    @nn.compact
+    def __call__(self, x):
+        x = jax.numpy.asarray(x, self.dtype)  # TODO: remove me?
+        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(
-    soft_majority_layer, hard_majority_layer, symbolic_majority_layer)
+    lambda dtype=jax.numpy.float32: SoftMajorityLayer(dtype),
+    lambda dtype=jax.numpy.float32: HardMajorityLayer(),
+    lambda dtype=jax.numpy.float32: SymbolicMajorityLayer()
+)

neurallogic/hard_or.py

@@ -61,7 +61,7 @@ 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)
+        x = jax.numpy.asarray(x, self.dtype) # TODO: remove me?
         return soft_or_layer(weights, x)
 
 

neurallogic/hard_xor.py

@@ -28,7 +28,7 @@ def soft_xor_neuron(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.0), xor, (0,))
+    x = jax.lax.reduce(x, jax.numpy.array(0, dtype=x.dtype), xor, (0,))
     return x
 
 
@@ -39,21 +39,21 @@ def hard_xor_neuron(w, x):
 
 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
+    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)
+        x = jax.numpy.asarray(x, self.dtype)  # TODO is this needed?
         return soft_xor_layer(weights, x)
 
 

neurallogic/harden.py

@@ -62,12 +62,16 @@ def harden(x: flax.core.FrozenDict):
     return harden(x.unfreeze())
 
 
+"""
 @dispatch
 def harden(*args):
     if len(args) == 1:
-        return harden(args[0])
+        print(f'args = {args} of type {type(args)}')
+        arg = args[0]
+        print(f'args[0] = {arg}')
+        return tuple(harden(arg))
     return tuple([harden(arg) for arg in args])
-
+"""
 
 @dispatch
 def map_keys_nested(f, d: dict) -> dict:

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/neural_logic_net.py

@@ -1,5 +1,6 @@
 from enum import Enum
 from flax import linen as nn
+from neurallogic import symbolic_generation
 
 NetType = Enum('NetType', ['Soft', 'Hard', 'Symbolic'])
 

neurallogic/symbolic_generation.py

@@ -100,6 +100,7 @@ def make_symbolic_flax_jaxpr(flax_layer, x):
     return jaxpr
 
 
+
 def eval_jaxpr(symbolic, jaxpr, consts, *args):
     '''Evaluates a jaxpr by interpreting it as Python code.
 
@@ -212,15 +213,15 @@ def eval_jaxpr_impl(jaxpr):
 def make_symbolic_jaxpr(func: typing.Callable, *args):
     return jax.make_jaxpr(lambda *args: func(*args))(*args)
 
-
-def eval_symbolic(symbolic_function, *args):
-    if hasattr(symbolic_function, 'literals'):
+# TODO: better name
+def eval_symbolic(jaxpr, *args):
+    if hasattr(jaxpr, 'literals'):
         return eval_jaxpr(
-            False, symbolic_function.jaxpr, symbolic_function.literals, *args
+            False, jaxpr.jaxpr, jaxpr.literals, *args
         )
-    return eval_jaxpr(False, symbolic_function.jaxpr, [], *args)
-
+    return eval_jaxpr(False, jaxpr.jaxpr, [], *args)
 
+# TODO: better name
 def symbolic_expression(jaxpr, *args):
     if hasattr(jaxpr, 'literals'):
         sym_expr = eval_jaxpr(True, jaxpr.jaxpr, jaxpr.literals, *args)

neurallogic/symbolic_operator.py

@@ -63,6 +63,11 @@ def symbolic_operator(operator: str, x: str, y: int):
     return symbolic_operator(operator, x, str(y))
 
 
+@dispatch
+def symbolic_operator(operator: str, x: tuple):
+    return symbolic_operator(operator, str(x))
+
+
 @dispatch
 def symbolic_operator(operator: str, x: list, y: numpy.ndarray):
     return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)

neurallogic/symbolic_primitives.py

@@ -21,83 +21,86 @@ def all_concrete_values(data):
     return True
 
 
-def symbolic_f(concrete_f: Callable, symbolic_f_name: str, *args, **kwargs):
+def symbolic(concrete_function: Callable, symbolic_function: str, *args, **kwargs):
     if all_concrete_values([*args]):
-        return concrete_f(*args, **kwargs)
+        # We can directly evaluate the function
+        return concrete_function(*args, **kwargs)
     else:
-        return symbolic_operator.symbolic_operator(symbolic_f_name, *args, **kwargs)
+        # We need to return a symbolic representation
+        return symbolic_operator.symbolic_operator(symbolic_function, *args, **kwargs)
 
 
 def symbolic_not(*args, **kwargs):
-    return symbolic_f(numpy.logical_not, 'numpy.logical_not', *args, **kwargs)
+    return symbolic(numpy.logical_not, 'numpy.logical_not', *args, **kwargs)
 
 
 def symbolic_eq(*args, **kwargs):
-    return symbolic_f(lax_reference.eq, 'lax_reference.eq', *args, **kwargs)
+    return symbolic(lax_reference.eq, 'lax_reference.eq', *args, **kwargs)
 
 
 def symbolic_ne(*args, **kwargs):
-    return symbolic_f(lax_reference.ne, 'lax_reference.ne', *args, **kwargs)
+    return symbolic(lax_reference.ne, 'lax_reference.ne', *args, **kwargs)
 
 
 def symbolic_le(*args, **kwargs):
-    return symbolic_f(lax_reference.le, 'lax_reference.le', *args, **kwargs)
+    return symbolic(lax_reference.le, 'lax_reference.le', *args, **kwargs)
 
 
 def symbolic_lt(*args, **kwargs):
-    return symbolic_f(lax_reference.lt, 'lax_reference.lt', *args, **kwargs)
+    return symbolic(lax_reference.lt, 'lax_reference.lt', *args, **kwargs)
 
 
 def symbolic_ge(*args, **kwargs):
-    return symbolic_f(lax_reference.ge, 'lax_reference.ge', *args, **kwargs)
+    return symbolic(lax_reference.ge, 'lax_reference.ge', *args, **kwargs)
 
 
 def symbolic_gt(*args, **kwargs):
-    return symbolic_f(lax_reference.gt, 'lax_reference.gt', *args, **kwargs)
+    return symbolic(lax_reference.gt, 'lax_reference.gt', *args, **kwargs)
 
 
 def symbolic_abs(*args, **kwargs):
-    return symbolic_f(lax_reference.abs, 'numpy.absolute', *args, **kwargs)
+    return symbolic(lax_reference.abs, 'numpy.absolute', *args, **kwargs)
 
 
 def symbolic_add(*args, **kwargs):
-    return symbolic_f(lax_reference.add, 'numpy.add', *args, **kwargs)
+    return symbolic(lax_reference.add, 'numpy.add', *args, **kwargs)
 
 
 def symbolic_sub(*args, **kwargs):
-    return symbolic_f(lax_reference.sub, 'numpy.subtract', *args, **kwargs)
+    return symbolic(lax_reference.sub, 'numpy.subtract', *args, **kwargs)
 
 
 def symbolic_mul(*args, **kwargs):
-    return symbolic_f(lax_reference.mul, 'numpy.multiply', *args, **kwargs)
+    return symbolic(lax_reference.mul, 'numpy.multiply', *args, **kwargs)
 
 
 def symbolic_div(*args, **kwargs):
-    return symbolic_f(lax_reference.div, 'lax_reference.div', *args, **kwargs)
+    return symbolic(lax_reference.div, 'lax_reference.div', *args, **kwargs)
 
 
 def symbolic_max(*args, **kwargs):
-    return symbolic_f(lax_reference.max, 'numpy.maximum', *args, **kwargs)
+    return symbolic(lax_reference.max, 'numpy.maximum', *args, **kwargs)
 
 
 def symbolic_min(*args, **kwargs):
-    return symbolic_f(lax_reference.min, 'numpy.minimum', *args, **kwargs)
+    return symbolic(lax_reference.min, 'numpy.minimum', *args, **kwargs)
 
 
 def symbolic_and(*args, **kwargs):
-    return symbolic_f(numpy.logical_and, 'numpy.logical_and', *args, **kwargs)
+    return symbolic(numpy.logical_and, 'numpy.logical_and', *args, **kwargs)
 
 
 def symbolic_or(*args, **kwargs):
-    return symbolic_f(numpy.logical_or, 'numpy.logical_or', *args, **kwargs)
+    return symbolic(numpy.logical_or, 'numpy.logical_or', *args, **kwargs)
 
 
 def symbolic_xor(*args, **kwargs):
-    return symbolic_f(numpy.logical_xor, 'numpy.logical_xor', *args, **kwargs)
+    return symbolic(numpy.logical_xor, 'numpy.logical_xor', *args, **kwargs)
 
 
 def symbolic_sum(*args, **kwargs):
-    return symbolic_f(lax_reference.sum, 'lax_reference.sum', *args, **kwargs)
+    # N.B. We pass the tuple directly because we're summing over all args
+    return symbolic(lax_reference.sum, 'lax_reference.sum', args, **kwargs)
 
 
 def symbolic_broadcast_in_dim(*args, **kwargs):

tests/test_hard_majority.py

@@ -2,6 +2,7 @@
 import jax
 
 from neurallogic import hard_majority, harden, symbolic_generation
+from tests import utils
 
 
 def test_majority_index():
@@ -126,9 +127,7 @@ def test_hard_majority_layer():
                      False, True, False, True, False, True], [False, True, False, True, False, True]])) == numpy.array([False, False, False]))
 
 
-def test_majority_layer():
-    soft, hard, symbolic = hard_majority.soft_majority_layer, hard_majority.hard_majority_layer, hard_majority.symbolic_majority_layer
-
+def test_layer():
     test_data = [
         [
             [[0.8, 0.1, 0.4], [1.0, 0.0, 0.3]],
@@ -151,16 +150,15 @@ def test_majority_layer():
             [0.3, 0.1, 0.2, 0.3, 0.0]
         ]
     ]
-    
+
     for input, expected in test_data:
-        input = jax.numpy.array(input)
-        expected = jax.numpy.array(expected)
-        soft_output = soft(input)
-        assert jax.numpy.array_equal(soft_output, expected)
-        hard_output = hard(harden.harden(input))
-        assert jax.numpy.array_equal(hard_output, harden.harden(expected))
-        jaxpr = symbolic_generation.make_symbolic_jaxpr(symbolic, harden.harden(input))
-        symbolic_output = symbolic_generation.symbolic_expression(jaxpr, harden.harden(input))
-        assert jax.numpy.array_equal(symbolic_output, harden.harden(expected))
+        def soft(input):
+            return hard_majority.soft_majority_layer(input)
+
+        def hard(input):
+            return hard_majority.hard_majority_layer(input)
+
+        utils.check_consistency(soft, hard, jax.numpy.array(expected), jax.numpy.array(input))
+
 
 # TODO: test training the hard majority layer