cleanup

Author: Z80coder

neurallogic/hard_and.py

@@ -4,7 +4,6 @@
 
 import numpy
 import jax
-from flax import errors
 from flax import linen as nn
 
 from neurallogic import neural_logic_net, sym_gen, symbolic_primitives
@@ -22,39 +21,10 @@ def soft_and_include(w: float, x: float) -> float:
     return jax.numpy.maximum(x, 1.0 - w)
 
 
-# TODO: do we need to jit here? should apply jit at the highest level of the architecture
-# TODO: may need to jit in unit tests, however
-"""
-@jax.jit
-def hard_and_include(w: bool, x: bool) -> bool:
-    print(f"hard_and_include: w={w}, x={x}")
-    # TODO: this works when the function is jitted, but not when it is not jitted
-    return x | ~w
-    #return x or not w
-"""
-
 
 def hard_and_include(w, x):
     return jax.numpy.logical_or(x, jax.numpy.logical_not(w))
 
-# def hard_and_include(w, x):
-#    return jax.numpy.logical_or(x, jax.numpy.logical_not(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 symbolic_and_include(w, x):
-#    symbolic_f = sym_gen.make_symbolic(hard_and_include, w, x)
-#    return sym_gen.eval_symbolic(symbolic_f, w, x)
 
 
 def soft_and_neuron(w, x):
@@ -67,49 +37,14 @@ def hard_and_neuron(w, x):
     return jax.lax.reduce(x, True, jax.lax.bitwise_and, [0])
 
 
-"""
-def hard_and_neuron(w, x):
-    x = jax.vmap(hard_and_include, 0, 0)(w, x)
-    return jax.lax.reduce(x, True, jax.numpy.logical_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]
-"""
-
-# def symbolic_and_layer(w, x):
-#    symbolic_hard_and_layer = sym_gen.make_symbolic(hard_and_layer)
-#    return sym_gen.eval_symbolic(symbolic_hard_and_layer, w, x)
-
-# 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)

neurallogic/sym_gen.py

@@ -4,7 +4,7 @@
 from jax import core
 from jax._src.util import safe_map
 import flax
-from neurallogic import symbolic_primitives, harden
+from neurallogic import symbolic_primitives
 from plum import dispatch
 import typing
 

tests/test_hard_and.py

@@ -9,7 +9,7 @@
 from neurallogic import hard_and, harden, neural_logic_net, primitives, sym_gen, symbolic_primitives
 
 
-def check_consistency(soft: typing.Callable, hard: typing.Callable,  symbolic: typing.Callable, expected, *args):
+def check_consistency(soft: typing.Callable, hard: typing.Callable, expected, *args):
     # Check that the soft function performs as expected
     assert numpy.allclose(soft(*args), expected)
 
@@ -18,18 +18,11 @@ def check_consistency(soft: typing.Callable, hard: typing.Callable,  symbolic: t
     hard_expected = harden.harden(expected)
     assert numpy.allclose(hard(*hard_args), hard_expected)
 
-    # Check that the symbolic function performs as expected
-    symbolic_f = sym_gen.make_symbolic_jaxpr(symbolic, *hard_args)
+    # Check that the jaxpr expression performs as expected
+    symbolic_f = sym_gen.make_symbolic_jaxpr(hard, *hard_args)
     assert numpy.allclose(sym_gen.eval_symbolic(
         symbolic_f, *hard_args), hard_expected)
 
-    # Check that the symbolic function, when evaluted with symbolic inputs, performs as expected
-    symbolic_input = sym_gen.make_symbolic(*hard_args)
-    symbolic_expression = sym_gen.symbolic_expression(
-        symbolic_f, *symbolic_input)
-    symbolic_output = sym_gen.eval_symbolic_expression(symbolic_expression)
-    assert numpy.allclose(symbolic_output, hard_expected)
-
 
 def test_include():
     test_data = [
@@ -44,7 +37,7 @@ def test_include():
     ]
     for input, expected in test_data:
         check_consistency(hard_and.soft_and_include, hard_and.hard_and_include,
-                          hard_and.hard_and_include, expected, input[0], input[1])
+                           expected, input[0], input[1])
 
 
 def test_neuron():
@@ -63,10 +56,7 @@ def soft(weights, input):
         def hard(weights, input):
             return hard_and.hard_and_neuron(weights, input)
 
-        def symbolic(weights, input):
-            return hard(weights, input)
-
-        check_consistency(soft, hard, symbolic, expected,
+        check_consistency(soft, hard, expected,
                           jax.numpy.array(weights), jax.numpy.array(input))
 
 
@@ -90,10 +80,8 @@ def soft(weights, input):
         def hard(weights, input):
             return hard_and.hard_and_layer(weights, input)
 
-        def symbolic(weights, input):
-            return hard(weights, input)
 
-        check_consistency(soft, hard, symbolic, jax.numpy.array(expected),
+        check_consistency(soft, hard, jax.numpy.array(expected),
                           jax.numpy.array(weights), jax.numpy.array(input))