refactor code

Author: Z80coder

neurallogic/harden.py

@@ -2,7 +2,8 @@
 import jax
 import numpy
 from plum import dispatch
-from neurallogic import symbolic_primitives
+
+from neurallogic import map_at_elements
 
 
 def harden_float(x: float) -> bool:
@@ -21,7 +22,7 @@ def harden(x: float):
 
 @dispatch
 def harden(x: list):
-    return symbolic_primitives.map_at_elements(x, harden_float)
+    return map_at_elements.map_at_elements(x, harden_float)
 
 
 @dispatch
@@ -39,7 +40,7 @@ def harden(x: dict):
     # Only harden parameters that explicitly represent bits
     def conditional_harden(k, v):
         if k.startswith("bit_"):
-            return symbolic_primitives.map_at_elements(v, harden)
+            return map_at_elements.map_at_elements(v, harden)
         elif isinstance(v, dict) or isinstance(v, flax.core.FrozenDict) or isinstance(v, list):
             return harden(v)
         return v

neurallogic/map_at_elements.py

@@ -0,0 +1,58 @@
+import typing
+
+import jax
+import numpy
+from plum import dispatch
+
+
+@dispatch
+def map_at_elements(x: str, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: bool, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: numpy.bool_, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: float, func: typing.Callable):
+    return func(x)
+
+
+@dispatch
+def map_at_elements(x: numpy.float32, 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]
+
+
+@dispatch
+def map_at_elements(x: numpy.ndarray, func: typing.Callable):
+    return numpy.array([map_at_elements(item, func) for item in x], dtype=object)
+
+
+@dispatch
+def map_at_elements(x: jax.numpy.ndarray, func: typing.Callable):
+    if x.ndim == 0:
+        return func(x.item())
+    return jax.numpy.array([map_at_elements(item, func) for item in x])
+
+
+@dispatch
+def map_at_elements(x: dict, func: typing.Callable):
+    return {k: map_at_elements(v, func) for k, v in x.items()}
+
+
+@dispatch
+def map_at_elements(x: tuple, func: typing.Callable):
+    return tuple(map_at_elements(list(x), func))
+

neurallogic/real_encoder.py

@@ -5,6 +5,7 @@
 
 from neurallogic import neural_logic_net, symbolic_generation
 
+
 def soft_real_encoder(t: float, x: float) -> float:
     eps = 0.0000001
     # x should be in [0, 1]
@@ -44,7 +45,8 @@ class SoftRealEncoderLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         thresholds_shape = (jax.numpy.shape(x)[-1], self.bits_per_real)
-        thresholds = self.param("thresholds", self.thresholds_init, thresholds_shape, self.dtype)
+        thresholds = self.param(
+            "thresholds", self.thresholds_init, thresholds_shape, self.dtype)
         x = jax.numpy.asarray(x, self.dtype)
         return soft_real_encoder_layer(thresholds, x)
 
@@ -55,7 +57,8 @@ class HardRealEncoderLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         thresholds_shape = (jax.numpy.shape(x)[-1], self.bits_per_real)
-        thresholds = self.param("thresholds", nn.initializers.constant(0.0), thresholds_shape)
+        thresholds = self.param(
+            "thresholds", nn.initializers.constant(0.0), thresholds_shape)
         return hard_real_encoder_layer(thresholds, x)
 
 

neurallogic/symbolic_generation.py

@@ -7,7 +7,7 @@
 from jax._src.util import safe_map
 from plum import dispatch
 
-from neurallogic import symbolic_primitives
+from neurallogic import symbolic_primitives, map_at_elements
 
 # Imports required for evaluating symbolic expressions with eval()
 import jax._src.lax_reference as lax_reference
@@ -74,7 +74,7 @@ def convert_to_numeric_params(flax_layer, param_names: str):
     if isinstance(actual_weights, list) or (
         isinstance(actual_weights, numpy.ndarray) and actual_weights.dtype == object
     ):
-        numeric_weights = symbolic_primitives.map_at_elements(
+        numeric_weights = map_at_elements.map_at_elements(
             actual_weights, lambda x: 0
         )
         numeric_weights = numpy.asarray(numeric_weights, dtype=numpy.int32)
@@ -87,7 +87,7 @@ def make_symbolic_flax_jaxpr(flax_layer, x):
     flax_layer, thresholds = convert_to_numeric_params(flax_layer, 'thresholds')
     # Convert input to dummy numeric input (if needed)
     if isinstance(x, list) or (isinstance(x, numpy.ndarray) and x.dtype == object):
-        x = symbolic_primitives.map_at_elements(x, lambda x: 0)
+        x = map_at_elements.map_at_elements(x, lambda x: 0)
         x = numpy.asarray(x, dtype=numpy.int32)
     # Make the jaxpr that corresponds to the flax layer
     jaxpr = make_symbolic_jaxpr(flax_layer, x)

neurallogic/symbolic_operator.py

@@ -0,0 +1,84 @@
+import jax
+import numpy
+from plum import dispatch
+
+
+@dispatch
+def symbolic_operator(operator: str, x: str) -> str:
+    return f'{operator}({x})'.replace('\'', '')
+
+
+@dispatch
+def symbolic_operator(operator: str, x: float, y: str):
+    return symbolic_operator(operator, str(x), y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: str, y: float):
+    return symbolic_operator(operator, x, str(y))
+
+
+@dispatch
+def symbolic_operator(operator: str, x: float, y: numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: numpy.ndarray, y: numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: str, y: str):
+    return f'{operator}({x}, {y})'.replace('\'', '')
+
+
+@dispatch
+def symbolic_operator(operator: str, x: numpy.ndarray, y: float):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: list, y: float):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: list, y: list):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: bool, y: str):
+    return symbolic_operator(operator, str(x), y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: str, y: numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: str, y: int):
+    return symbolic_operator(operator, x, str(y))
+
+
+@dispatch
+def symbolic_operator(operator: str, x: list, y: numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: numpy.ndarray, y: jax.numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x, y)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: numpy.ndarray):
+    return numpy.vectorize(symbolic_operator, otypes=[object])(operator, x)
+
+
+@dispatch
+def symbolic_operator(operator: str, x: list):
+    return symbolic_operator(operator, numpy.array(x))
+