all tests passing locally

Author: Z80coder

neurallogic/symbolic_generation.py

@@ -121,8 +121,8 @@ def eval_jaxpr(symbolic, jaxpr, consts, *args):
     out : tuple
         The result of evaluating the jaxpr.
     '''
-    if symbolic:
-        numpy.set_printoptions(threshold=sys.maxsize)
+    #if symbolic:
+    #    numpy.set_printoptions(threshold=sys.maxsize)
 
     # Mapping from variable -> value
     env = {}
@@ -313,7 +313,7 @@ def eval_symbolic_expression(x: numpy.ndarray):
     return numpy.vectorize(eval_symbolic_expression)(x)
 
 
-#@dispatch
-#def eval_symbolic_expression(x: list):
-#    return numpy.vectorize(eval)(x)
+@dispatch
+def eval_symbolic_expression(x: list):
+    return numpy.vectorize(eval)(x)
 

neurallogic/symbolic_primitives.py

@@ -8,6 +8,7 @@
 
 # TODO: remove me?
 def convert_element_type(x, dtype):
+    return x
     if dtype == numpy.int32 or dtype == numpy.int64:
         #dtype = 'int'
         return x
@@ -199,6 +200,7 @@ def symbolic_operator(operator: str, x: list):
 
 
 # TODO: remove infix_operator?
+"""
 @dispatch
 def symbolic_infix_operator(operator: str, a: str, b: str) -> str:
     return f'{a} {operator} {b}'.replace('\'', '')
@@ -248,7 +250,7 @@ def symbolic_infix_operator(operator: str, a: list, b: list):
 @dispatch
 def symbolic_infix_operator(operator: str, a: bool, b: str):
     return symbolic_infix_operator(operator, str(a), b)
-
+"""
 
 def all_concrete_values(data):
     if isinstance(data, str):
@@ -275,39 +277,34 @@ def symbolic_eq(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.eq(*args, **kwargs)
     else:
-        #return '(' + symbolic_infix_operator('==', *args, **kwargs) + ')'
         return symbolic_operator('lax_reference.eq', *args, **kwargs)
 
 
 def symbolic_ne(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.ne(*args, **kwargs)
     else:
-        #return '(' + symbolic_infix_operator('!=', *args, **kwargs) + ')'
         return symbolic_operator('lax_reference.ne', *args, **kwargs)
 
 
 def symbolic_le(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.le(*args, **kwargs)
     else:
-        #return '(' + symbolic_infix_operator('<=', *args, **kwargs) + ')'
         return symbolic_operator('lax_reference.le', *args, **kwargs)
 
 
 def symbolic_lt(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.lt(*args, **kwargs)
     else:
-        #return '(' + symbolic_infix_operator('<', *args, **kwargs) + ')'
         return symbolic_operator('lax_reference.lt', *args, **kwargs)
 
 
 def symbolic_gt(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.gt(*args, **kwargs)
     else:
-        #return symbolic_infix_operator('>', *args, **kwargs)
         return symbolic_operator('lax_reference.gt', *args, **kwargs)
 
 
@@ -393,7 +390,6 @@ def symbolic_and(*args, **kwargs):
     if all_concrete_values([*args]):
         return numpy.logical_and(*args, **kwargs)
     else:
-        #return symbolic_infix_operator('and', *args, **kwargs)
         return symbolic_operator('numpy.logical_and', *args, **kwargs)
 
 
@@ -402,33 +398,27 @@ def symbolic_or(*args, **kwargs):
     if all_concrete_values([*args]):
         return numpy.logical_or(*args, **kwargs)
     else:
-        #return '(' + symbolic_infix_operator('or', *args, **kwargs) + ')'
         return symbolic_operator('numpy.logical_or', *args, **kwargs)
 
 
 def symbolic_xor(*args, **kwargs):
     if all_concrete_values([*args]):
         return numpy.logical_xor(*args, **kwargs)
     else:
-        #return symbolic_infix_operator('^', *args, **kwargs)
         return symbolic_operator('numpy.logical_xor', *args, **kwargs)
 
 
 def symbolic_sum(*args, **kwargs):
     if all_concrete_values([*args]):
         return lax_reference.sum(*args, **kwargs)
     else:
-        #return symbolic_infix_operator('+', *args, **kwargs)
         return symbolic_operator('lax_reference.sum', *args, **kwargs)
 
 
 def symbolic_broadcast_in_dim(*args, **kwargs):
-    assert len(args) == 1
-    arg = args[0]
-    if isinstance(args, (list, tuple)):
-        # reference implementation demands a numpy array
-        arg = numpy.array(arg)
-    return lax_reference.broadcast_in_dim(arg, **kwargs)
+    # broadcast_in_dim requires numpy arrays not lists
+    args = tuple([numpy.array(arg) if isinstance(arg, list) else arg for arg in args])
+    return lax_reference.broadcast_in_dim(*args, **kwargs)
 
 
 def symbolic_reshape(*args, **kwargs):