clean-up handling of type conversions, and mapping of elements in arbitrary nested data structures

Author: Z80coder

neurallogic/sym_gen.py

@@ -129,8 +129,7 @@ def eval_jaxpr_concrete(jaxpr, *args):
 
 
 def eval_jaxpr_symbolic(jaxpr, *args):
-    # Convert the literals to symbolic literals
-    symbolic_jaxpr_literals = symbolic_primitives.to_boolean_symbolic_values(
-        jaxpr.literals)
+    symbolic_jaxpr_literals = safe_map(lambda x: numpy.array(x, dtype=object), jaxpr.literals)
+    symbolic_jaxpr_literals = symbolic_primitives.to_boolean_symbolic_values(symbolic_jaxpr_literals)
     return eval_jaxpr(True, jaxpr.jaxpr, symbolic_jaxpr_literals, *args)
 

neurallogic/symbolic_primitives.py

@@ -1,71 +1,72 @@
 import numpy
 from plum import dispatch
+import typing
 import jax
 import jax._src.lax_reference as lax_reference
-from neurallogic import primitives
-
-
-def to_boolean_value_string(x):
-    if isinstance(x, bool):
-        # x is a bool
-        return 'True' if x else 'False'
-    elif x == 1.0 or x == 0.0:
-        # x is a float
-        return 'True' if x == 1.0 else 'False'
-    elif isinstance(x, str) and (x == '1' or x == '0'):
-        # x is a string representing an integer
-        return 'True' if x == '1' else 'False'
-    elif isinstance(x, str) and (x == '1.0' or x == '0.0'):
-        # x is a string representing a float
-        return 'True' if x == '1.0' else 'False'
-    elif isinstance(x, str) and (x == 'True' or x == 'False'):
-        # x is a string representing a boolean
+import jaxlib
+
+def convert_iterable_type(x: list, new_type):
+    if new_type == list:
         return x
-    elif isinstance(x, numpy.ndarray) or isinstance(x, jax.numpy.ndarray) or isinstance(x, list) or isinstance(x, tuple):
-        # We only operate on scalars
-        raise ValueError(
-            f"to_boolean_value_string only operates on scalars, but got {x}")
+    elif new_type == numpy.ndarray:
+        return numpy.array(x, dtype=object)
+    elif new_type == jax.numpy.ndarray:
+        return jax.numpy.array(x, dtype=object)
+    elif new_type == jaxlib.xla_extension.DeviceArray:
+        return jax.numpy.array(x, dtype=object)
     else:
-        # x is not interpretable as a boolean
-        return str(x)
+        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'
 
-def to_boolean_symbolic_values_impl(x):
-    """Converts an arbitrary vector of arbitrary values to a list where
-    every boolean-interpretable value gets converted to the strings "True" or "False".
+@dispatch
+def to_boolean_value_string(x: int):
+    return 'True' if x == 1.0 else 'False'
 
-    Args:
-        x: The vector of values to convert (or can be a single value in the degenerate case)
+@dispatch
+def to_boolean_value_string(x: float):
+    return 'True' if x == 1.0 else 'False'
 
-    Returns:
-        A list representation of the input, where boolean-interpretable
-        values are converted to "True" or "False".
-    """
-    if isinstance(x, numpy.ndarray) or isinstance(x, jax.numpy.ndarray) or isinstance(x, tuple):
-        return to_boolean_symbolic_values_impl(x.tolist())
-    elif isinstance(x, list):
-        return [to_boolean_symbolic_values_impl(y) for y in x]
+@dispatch
+def to_boolean_value_string(x: str):
+    if x == '1' or x == '1.0' or x =='True':
+        return 'True'
+    elif x == '0' or x == '0.0' or x =='False':
+        return 'False'
     else:
-        return to_boolean_value_string(x)
+        return x
 
 
 def to_boolean_symbolic_values(x):
-    """Converts an arbitrary vector of arbitrary values to a numpy array where
-    every boolean-interpretable value gets converted to the strings "True" or "False".
-
-    Args:
-        x: The vector of values to convert (or can be a single value in the degenerate case)
-
-    Returns:
-        A numpy array representation of the input, where boolean-interpretable
-        values are converted to "True" or "False".
-    """
-    x = to_boolean_symbolic_values_impl(x)
-    if isinstance(x, list):
-        x = numpy.array(x, dtype=object)
-    else:
-        x = numpy.array([x], dtype=object)
-    return x
+    return map_at_elements(x, to_boolean_value_string)
 
 
 @dispatch
@@ -167,42 +168,12 @@ def symbolic_broadcast_in_dim(*args, **kwargs):
     return lax_reference.broadcast_in_dim(*args, **kwargs)
 
 
-def is_iterable(obj):
-    try:
-        iter(obj)
-        return True
-    except TypeError:
-        return False
-
-# TODO: unify this way of walking a nested iterable with the code above
-def apply_func_to_nested_impl(iterable, func):
-    if isinstance(iterable, (numpy.ndarray, jax.numpy.ndarray)):
-        iterable = iterable.tolist()
-    if is_iterable(iterable):
-        transformed = []
-        for item in iterable:
-            if isinstance(item, list):
-                transformed.append(apply_func_to_nested_impl(item, func))
-            else:
-                transformed.append(func(item))
-        return transformed
-    else:
-        return func(iterable)
-
-def apply_func_to_nested(iterable, func):
-    iterable_type = type(iterable)
-    r = apply_func_to_nested_impl(iterable, func)
-    if iterable_type == numpy.ndarray:
-        r = numpy.array(r, dtype=object)
-    assert type(r) == iterable_type
-    return r
-
 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 apply_func_to_nested(x, convert)
+    return map_at_elements(x, convert)
 
 
 # TODO: add a test for this

tests/test_sym_gen.py

@@ -53,7 +53,7 @@ def test_sym_gen():
     # -- TEST 2: Compare the standard evaluation of the network with the non-standard symbolic evaluation of the jaxpr
     # Convert the hard input to a symbolic input
     symbolic_mock_input = symbolic_primitives.to_boolean_symbolic_values(
-        hard_mock_input)
+        numpy.array(hard_mock_input, dtype=object))
     # Symbolically evaluate the jaxpr with the symbolic input
     eval_symbolic_output = sym_gen.eval_jaxpr_symbolic(
         jaxpr, symbolic_mock_input)

tests/test_symbolic_primitives.py

@@ -141,42 +141,42 @@ def test_to_boolean_value_string():
 
 def test_to_boolean_symbolic_values():
     output = symbolic_primitives.to_boolean_symbolic_values([1, 1])
-    expected = numpy.array(["True", "True"])
+    expected = ["True", "True"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([0, 0])
-    expected = numpy.array(["False", "False"])
+    expected = ["False", "False"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([True, False])
-    expected = numpy.array(["True", "False"])
+    expected = ["True", "False"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([False, True])
-    expected = numpy.array(["False", "True"])
+    expected = ["False", "True"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([1.0, 1.0])
-    expected = numpy.array(["True", "True"])
+    expected = ["True", "True"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([0.0, 0.0])
-    expected = numpy.array(["False", "False"])
+    expected = ["False", "False"]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([[1, 1], [1, 1]])
-    expected = numpy.array([["True", "True"], ["True", "True"]])
+    expected = [["True", "True"], ["True", "True"]]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values([[0, 0], [0, 0]])
-    expected = numpy.array([["False", "False"], ["False", "False"]])
+    expected = [["False", "False"], ["False", "False"]]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values(
         [[True, False], [False, True]])
-    expected = numpy.array([["True", "False"], ["False", "True"]])
+    expected = [["True", "False"], ["False", "True"]]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values(
         [[[1, 0, 1], [1, 0, 1]], [[1, 0, 0], [1, 0, 0]]])
-    expected = numpy.array([[["True", "False", "True"], ["True", "False", "True"]], [
-                           ["True", "False", "False"], ["True", "False", "False"]]])
+    expected = [[["True", "False", "True"], ["True", "False", "True"]], [
+        ["True", "False", "False"], ["True", "False", "False"]]]
     assert numpy.array_equal(output, expected)
     output = symbolic_primitives.to_boolean_symbolic_values(
         [[[1, "f", 1], [1, "g", 1]], [[1, "h", 0], [1, "f", 0]]])
-    expected = numpy.array([[["True", "f", "True"], ["True", "g", "True"]], [
-                           ["True", "h", "False"], ["True", "f", "False"]]])
+    expected = [[["True", "f", "True"], ["True", "g", "True"]], [
+        ["True", "h", "False"], ["True", "f", "False"]]]
     assert numpy.array_equal(output, expected)