fix bug

Author: Z80coder

neurallogic/symbolic_primitives.py

@@ -5,6 +5,7 @@
 import jax._src.lax_reference as lax_reference
 import jaxlib
 
+
 def convert_iterable_type(x: list, new_type):
     if new_type == list:
         return x
@@ -15,51 +16,62 @@ def convert_iterable_type(x: list, new_type):
     elif new_type == jaxlib.xla_extension.DeviceArray:
         return jax.numpy.array(x, dtype=object)
     else:
-        raise NotImplementedError(f"Cannot convert type {type(x)} to type {new_type}")
+        raise NotImplementedError(
+            f"Cannot convert type {type(x)} to type {new_type}")
+
 
 @dispatch
 def map_at_elements(x: list, func: typing.Callable):
     return convert_iterable_type([map_at_elements(item, func) for item in x], type(x))
 
+
 @dispatch
 def map_at_elements(x: numpy.ndarray, func: typing.Callable):
     return convert_iterable_type([map_at_elements(item, func) for item in x], type(x))
 
+
 @dispatch
 def map_at_elements(x: jax.numpy.ndarray, func: typing.Callable):
     if x.ndim == 0:
         return func(x.item())
     return convert_iterable_type([map_at_elements(item, func) for item in x], type(x))
 
+
 @dispatch
 def map_at_elements(x: str, func: typing.Callable):
     return func(x)
 
+
 @dispatch
 def map_at_elements(x, func: typing.Callable):
     return func(x)
 
+
 @dispatch
 def to_boolean_value_string(x: bool):
     return 'True' if x else 'False'
 
+
 @dispatch
 def to_boolean_value_string(x: numpy.bool_):
     return 'True' if x else 'False'
 
+
 @dispatch
 def to_boolean_value_string(x: int):
     return 'True' if x == 1.0 else 'False'
 
+
 @dispatch
 def to_boolean_value_string(x: float):
     return 'True' if x == 1.0 else 'False'
 
+
 @dispatch
 def to_boolean_value_string(x: str):
-    if x == '1' or x == '1.0' or x =='True':
+    if x == '1' or x == '1.0' or x == 'True':
         return 'True'
-    elif x == '0' or x == '0.0' or x =='False':
+    elif x == '0' or x == '0.0' or x == 'False':
         return 'False'
     else:
         return x
@@ -86,7 +98,6 @@ def unary_operator(operator: str, x: list):
 
 @dispatch
 def binary_infix_operator(operator: str, a: str, b: str, bracket: bool = False) -> str:
-    # We need to specify bracket because Python cannot evaluate expressions with too many nested parantheses
     if bracket:
         return f"({a}) {operator} ({b})"
     return f"{a} {operator} {b}"
@@ -160,17 +171,17 @@ def symbolic_sum(*args, **kwargs):
     else:
         return binary_infix_operator("+", *args, **kwargs)
 
-# Uses the lax reference implementation of broadcast_in_dim to
-# implement a symbolic version of broadcast_in_dim
-
 
 def symbolic_broadcast_in_dim(*args, **kwargs):
+    # Uses the lax reference implementation of broadcast_in_dim to
+    # implement a symbolic version of broadcast_in_dim
     return lax_reference.broadcast_in_dim(*args, **kwargs)
 
 
 def symbolic_convert_element_type_impl(x, dtype):
     if dtype == numpy.int32 or dtype == numpy.int64:
         dtype = "int"
+
     def convert(x):
         return f"{dtype}({x})"
     return map_at_elements(x, convert)
@@ -187,13 +198,13 @@ def symbolic_convert_element_type(*args, **kwargs):
         return symbolic_convert_element_type_impl(*args, dtype=kwargs['new_dtype'])
 
 
-# This function is a hack to get around the fact that JAX doesn't
-# support symbolic reduction operations. It takes a symbolic reduction
-# operation and a symbolic initial value and returns a function that
-# performs the reduction operation on a numpy array.
 
 
 def make_symbolic_reducer(py_binop, init_val):
+    # This function is a hack to get around the fact that JAX doesn't
+    # support symbolic reduction operations. It takes a symbolic reduction
+    # operation and a symbolic initial value and returns a function that
+    # performs the reduction operation on a numpy array.
     def reducer(operand, axis):
         # axis=None means we are reducing over all axes of the operand.
         axis = range(numpy.ndim(operand)) if axis is None else axis

tests/test_sym_gen.py

@@ -7,7 +7,6 @@
 
 def nln(type, x, width):
     x = hard_or.or_layer(type)(width)(x)
-    # if not_layer has size "width" then this fails. why?
     x = hard_not.not_layer(type)(10)(x)
     x = primitives.nl_ravel(type)(x)
     x = harden_layer.harden_layer(type)(x)
@@ -26,7 +25,7 @@ def test_sym_gen():
         test_ds["image"], (test_ds["image"].shape[0], -1))
 
     # Define width of network
-    width = 4
+    width = 10
     # Define the neural logic net
     soft, hard, _ = neural_logic_net.net(lambda type, x: nln(type, x, width))
     # Initialize a random number generator