wip: adding a hard_bit layer

Author: Z80coder

neurallogic/hard_bit.py

@@ -0,0 +1,29 @@
+import jax
+
+
+def soft_bit(t: float, x: float) -> float:
+    # x should be in [0, 1]
+    t = jax.numpy.clip(t, 0.0, 1.0)
+    return jax.numpy.where(
+        x == t,
+        0.5,
+        jax.numpy.where(
+            x < t,
+            (1.0 / (2.0 * t)) * x,
+            (1.0 / (2.0 * (1.0 - t))) * (x + 1.0 - 2.0 * t),
+        ),
+    )
+
+
+def hard_bit(t: float, x: float) -> bool:
+    # t and x must be floats
+    return jax.numpy.where(soft_bit(t, x) > 0.5, True, False)
+
+
+soft_bit_neuron = jax.vmap(soft_bit, in_axes=(0, None))
+
+hard_bit_neuron = jax.vmap(hard_bit, in_axes=(0, None))
+
+soft_bit_layer = jax.vmap(soft_bit_neuron, (0, 0), 0)
+
+hard_bit_layer = jax.vmap(hard_bit_neuron, (0, 0), 0)

neurallogic/hard_not.py

@@ -28,12 +28,9 @@ def hard_not(w: bool, x: bool) -> bool:
 hard_not_neuron = jax.vmap(hard_not, 0, 0)
 
 
+soft_not_layer = jax.vmap(soft_not_neuron, (0, 0), 0)
 
-soft_not_layer = jax.vmap(soft_not_neuron, (0, None), 0)
-
-hard_not_layer = jax.vmap(hard_not_neuron, (0, None), 0)
-
-
+hard_not_layer = jax.vmap(hard_not_neuron, (0, 0), 0)
 
 
 class SoftNotLayer(nn.Module):
@@ -44,8 +41,7 @@ class SoftNotLayer(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_not_layer(weights, x)
 
@@ -56,8 +52,7 @@ class HardNotLayer(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_not_layer(weights, x)
 
 
@@ -73,7 +68,12 @@ def __call__(self, x):
 
 not_layer = neural_logic_net.select(
     lambda layer_size, weights_init=nn.initializers.uniform(
-        1.0), dtype=jax.numpy.float32: SoftNotLayer(layer_size, weights_init, dtype),
+        1.0
+    ), dtype=jax.numpy.float32: SoftNotLayer(layer_size, weights_init, dtype),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: HardNotLayer(layer_size),
     lambda layer_size, weights_init=nn.initializers.uniform(
-        1.0), dtype=jax.numpy.float32: HardNotLayer(layer_size),
-    lambda layer_size, weights_init=nn.initializers.uniform(1.0), dtype=jax.numpy.float32: SymbolicNotLayer(layer_size))
+        1.0
+    ), dtype=jax.numpy.float32: SymbolicNotLayer(layer_size),
+)

neurallogic/harden.py

@@ -8,10 +8,13 @@
 def harden_float(x: float) -> bool:
     return x > 0.5
 
+
 harden_array = jax.vmap(harden_float, 0, 0)
 
 @dispatch
 def harden(x: float):
+    if numpy.isnan(x):
+        return x
     return harden_float(x)
 
 @dispatch

neurallogic/symbolic_generation.py

@@ -1,43 +1,51 @@
 import numpy
 import jax
-import jax._src.lax_reference as lax_reference
 from jax import core
 from jax._src.util import safe_map
 import flax
 from neurallogic import symbolic_primitives
 from plum import dispatch
 import typing
-from typing import (Any, Mapping)
-
-# TODO: rename this file to symbolic.py
+from typing import Any, Mapping
 
 
 def symbolic_bind(prim, *args, **params):
     # print("\nprimitive: ", prim.name)
     # print("args: ", args)
     # print("params: ", params)
     symbolic_outvals = {
-        'broadcast_in_dim': symbolic_primitives.symbolic_broadcast_in_dim,
-        'reshape': lax_reference.reshape,
-        'convert_element_type': symbolic_primitives.symbolic_convert_element_type,
-        'and': symbolic_primitives.symbolic_and,
-        'or': symbolic_primitives.symbolic_or,
-        'xor': symbolic_primitives.symbolic_xor,
-        'not': symbolic_primitives.symbolic_not,
-        'ne': symbolic_primitives.symbolic_ne,
-        'gt': symbolic_primitives.symbolic_gt,
-        'reduce_and': symbolic_primitives.symbolic_reduce_and,
-        'reduce_or': symbolic_primitives.symbolic_reduce_or,
-        'reduce_sum': symbolic_primitives.symbolic_reduce_sum,
+        "broadcast_in_dim": symbolic_primitives.symbolic_broadcast_in_dim,
+        "reshape": symbolic_primitives.symbolic_reshape,
+        "transpose": symbolic_primitives.symbolic_transpose,
+        "convert_element_type": symbolic_primitives.symbolic_convert_element_type,
+        "eq": symbolic_primitives.symbolic_eq,
+        "ne": symbolic_primitives.symbolic_ne,
+        "le": symbolic_primitives.symbolic_le,
+        "lt": symbolic_primitives.symbolic_lt,
+        "gt": symbolic_primitives.symbolic_gt,
+        "add": symbolic_primitives.symbolic_add,
+        "sub": symbolic_primitives.symbolic_sub,
+        "mul": symbolic_primitives.symbolic_mul,
+        "div": symbolic_primitives.symbolic_div,
+        "max": symbolic_primitives.symbolic_max,
+        "min": symbolic_primitives.symbolic_min,
+        "and": symbolic_primitives.symbolic_and,
+        "or": symbolic_primitives.symbolic_or,
+        "xor": symbolic_primitives.symbolic_xor,
+        "not": symbolic_primitives.symbolic_not,
+        "reduce_and": symbolic_primitives.symbolic_reduce_and,
+        "reduce_or": symbolic_primitives.symbolic_reduce_or,
+        "reduce_sum": symbolic_primitives.symbolic_reduce_sum,
+        "select_n": symbolic_primitives.symbolic_select_n,
     }[prim.name](*args, **params)
     return symbolic_outvals
 
+
 def scope_put_variable(self, col: str, name: str, value: Any):
     variables = self._collection(col)
 
     def put(target, key, val):
-        if (key in target and isinstance(target[key], dict) and
-                isinstance(val, Mapping)):
+        if key in target and isinstance(target[key], dict) and isinstance(val, Mapping):
             for k, v in val.items():
                 put(target[key], k, v)
         else:
@@ -50,11 +58,16 @@ def put_variable(self, col: str, name: str, value: Any):
     self.scope._variables = self.scope.variables().unfreeze()
     scope_put_variable(self.scope, col, name, value)
 
+
 def make_symbolic_flax_jaxpr(flax_layer, x):
     actual_weights = flax_layer.get_variable("params", "weights")
     # Convert actual weights to dummy numeric weights (if needed)
-    if isinstance(actual_weights, list) or (isinstance(actual_weights, numpy.ndarray) and actual_weights.dtype == object):
-        numeric_weights = symbolic_primitives.map_at_elements(actual_weights, lambda x: 0)
+    if isinstance(actual_weights, list) or (
+        isinstance(actual_weights, numpy.ndarray) and actual_weights.dtype == object
+    ):
+        numeric_weights = symbolic_primitives.map_at_elements(
+            actual_weights, lambda x: 0
+        )
         numeric_weights = numpy.asarray(numeric_weights, dtype=numpy.int32)
         put_variable(flax_layer, "params", "weights", numeric_weights)
     # Convert input to dummy numeric input (if needed)
@@ -130,33 +143,40 @@ def eval_jaxpr_impl(jaxpr):
             symbolic_invals = safe_map(symbolic_read, eqn.invars)
             prim = eqn.primitive
             if type(prim) is jax.core.CallPrimitive:
-                call_jaxpr = eqn.params['call_jaxpr']
+                call_jaxpr = eqn.params["call_jaxpr"]
                 if not symbolic:
                     safe_map(write, call_jaxpr.invars, map(read, eqn.invars))
                 try:
-                    safe_map(symbolic_write, call_jaxpr.invars,
-                             map(symbolic_read, eqn.invars))
+                    safe_map(
+                        symbolic_write,
+                        call_jaxpr.invars,
+                        map(symbolic_read, eqn.invars),
+                    )
                 except:
                     pass
                 eval_jaxpr_impl(call_jaxpr)
                 if not symbolic:
                     safe_map(write, eqn.outvars, map(read, call_jaxpr.outvars))
-                safe_map(symbolic_write, eqn.outvars, map(
-                    symbolic_read, call_jaxpr.outvars))
+                safe_map(
+                    symbolic_write, eqn.outvars, map(symbolic_read, call_jaxpr.outvars)
+                )
             else:
                 if not symbolic:
                     outvals = prim.bind(*invals, **eqn.params)
-                symbolic_outvals = symbolic_bind(
-                    prim, *symbolic_invals, **eqn.params)
+                symbolic_outvals = symbolic_bind(prim, *symbolic_invals, **eqn.params)
                 # Primitives may return multiple outputs or not
                 if not prim.multiple_results:
                     if not symbolic:
                         outvals = [outvals]
                     symbolic_outvals = [symbolic_outvals]
                 if not symbolic:
                     # Check that the concrete and symbolic values are equal
-                    assert numpy.array_equal(
-                        numpy.array(outvals), symbolic_outvals)
+                    #print(
+                    #    f"outvals: {outvals} and symbolic_outvals: {symbolic_outvals}"
+                    #)
+                    assert numpy.allclose(
+                        numpy.array(outvals), symbolic_outvals, equal_nan=True
+                    )
                 # Write the results of the primitive into the environment
                 if not symbolic:
                     safe_map(write, eqn.outvars, outvals)
@@ -169,24 +189,36 @@ def eval_jaxpr_impl(jaxpr):
     else:
         return safe_map(symbolic_read, jaxpr.outvars)[0]
 
+
 # TODO: parameterise these functions by the element conversion function
 
 # TODO: use union types to consolidate these functions
 @dispatch
 def make_symbolic(x: dict):
-    return symbolic_primitives.map_at_elements(x, symbolic_primitives.to_boolean_value_string)
+    return symbolic_primitives.map_at_elements(
+        x, symbolic_primitives.to_boolean_value_string
+    )
+
 
 @dispatch
 def make_symbolic(x: list):
-    return symbolic_primitives.map_at_elements(x, symbolic_primitives.to_boolean_value_string)
+    return symbolic_primitives.map_at_elements(
+        x, symbolic_primitives.to_boolean_value_string
+    )
+
 
 @dispatch
 def make_symbolic(x: numpy.ndarray):
-    return symbolic_primitives.map_at_elements(x, symbolic_primitives.to_boolean_value_string)
+    return symbolic_primitives.map_at_elements(
+        x, symbolic_primitives.to_boolean_value_string
+    )
+
 
 @dispatch
 def make_symbolic(x: jax.numpy.ndarray):
-    return symbolic_primitives.map_at_elements(convert_jax_to_numpy_arrays(x), symbolic_primitives.to_boolean_value_string)
+    return symbolic_primitives.map_at_elements(
+        convert_jax_to_numpy_arrays(x), symbolic_primitives.to_boolean_value_string
+    )
 
 
 @dispatch
@@ -214,25 +246,28 @@ def make_symbolic(x: flax.core.FrozenDict):
     x = convert_jax_to_numpy_arrays(x.unfreeze())
     return flax.core.FrozenDict(make_symbolic(x))
 
+
 @dispatch
 def make_symbolic(*args):
     return tuple([make_symbolic(arg) for arg in args])
 
+
 @dispatch
 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'):
-        return eval_jaxpr(False, symbolic_function.jaxpr, symbolic_function.literals, *args)
+    if hasattr(symbolic_function, "literals"):
+        return eval_jaxpr(
+            False, symbolic_function.jaxpr, symbolic_function.literals, *args
+        )
     return eval_jaxpr(False, symbolic_function.jaxpr, [], *args)
 
 
 def symbolic_expression(jaxpr, *args):
-    if hasattr(jaxpr, 'literals'):
-        sym_expr = eval_jaxpr(True, jaxpr.jaxpr,
-                              jaxpr.literals, *args)
+    if hasattr(jaxpr, "literals"):
+        sym_expr = eval_jaxpr(True, jaxpr.jaxpr, jaxpr.literals, *args)
     else:
         sym_expr = eval_jaxpr(True, jaxpr.jaxpr, [], *args)
     return sym_expr