Merge pull request #56 from github/multiple-backends

Document incompleteness of select_n

Author: Z80coder

neurallogic/real_encoder.py

@@ -5,22 +5,22 @@
 
 from neurallogic import neural_logic_net, symbolic_generation
 
-# TODO: perhaps this can be simplified with a simple multiplication?
 # TODO: implement a soft_real_decoder that can perhaps replace the port count approach
 
+
 def soft_real_encoder(t: float, x: float) -> float:
     eps = 0.0000001
     # x should be in [0, 1]
-    t = jax.numpy.clip(t, 0.0, 1.0)
+    t = jax.numpy.clip(t, 0, 1)
     return jax.numpy.where(
         jax.numpy.isclose(t, x),
         0.5,
         # t != x
         jax.numpy.where(
             x < t,
-            (1.0 / (2.0 * t + eps)) * x,
+            (x / (2 * t + eps)),
             # x > t
-            (1.0 / (2.0 * (1.0 - t) + eps)) * (x + 1.0 - 2.0 * t)
+            (x + 1 - 2 * t) / (2 * (1 - t) + eps)
         )
     )
 

neurallogic/symbolic_generation.py

@@ -28,13 +28,17 @@ def symbolic_bind(prim, *args, **params):
         'lt': symbolic_primitives.symbolic_lt,
         'ge': symbolic_primitives.symbolic_ge,
         'gt': symbolic_primitives.symbolic_gt,
-        'abs': symbolic_primitives.symbolic_abs,
         'add': symbolic_primitives.symbolic_add,
         'sub': symbolic_primitives.symbolic_sub,
         'mul': symbolic_primitives.symbolic_mul,
         'div': symbolic_primitives.symbolic_div,
+        'tan': symbolic_primitives.symbolic_tan,
         'max': symbolic_primitives.symbolic_max,
         'min': symbolic_primitives.symbolic_min,
+        'abs': symbolic_primitives.symbolic_abs,
+        'round': symbolic_primitives.symbolic_round,
+        'floor': symbolic_primitives.symbolic_floor,
+        'ceil': symbolic_primitives.symbolic_ceil,
         'and': symbolic_primitives.symbolic_and,
         'or': symbolic_primitives.symbolic_or,
         'xor': symbolic_primitives.symbolic_xor,
@@ -93,10 +97,11 @@ def make_symbolic_flax_jaxpr(flax_layer, x):
         x = numpy.asarray(x, dtype=numpy.int32)
     # Make the jaxpr that corresponds to the flax layer
     jaxpr = make_symbolic_jaxpr(flax_layer, x)
-    # Make a list of bit_weights and thresholds but only include each if they are not None
-    bit_weights_and_thresholds = [x for x in [bit_weights, thresholds] if x is not None]
-    # Replace the dummy numeric weights with the actual weights in the jaxpr
-    jaxpr.consts = bit_weights_and_thresholds
+    if hasattr(jaxpr, '_consts'):
+        # Make a list of bit_weights and thresholds but only include each if they are not None
+        bit_weights_and_thresholds = [x for x in [bit_weights, thresholds] if x is not None]
+        # Replace the dummy numeric weights with the actual weights in the jaxpr
+        jaxpr.__setattr__('_consts', bit_weights_and_thresholds)
     return jaxpr
 
 
@@ -190,10 +195,9 @@ def eval_jaxpr_impl(jaxpr):
                         outvals = [outvals]
                     symbolic_outvals = [symbolic_outvals]
                 if not symbolic:
-                    # Check that the concrete and symbolic values are equal
-                    # print(
-                    #    f'outvals: {outvals} and symbolic_outvals: {symbolic_outvals}'
-                    # )
+                    # Always check that the symbolic binding generates the same values as the
+                    # standard jax binding in order to detect bugs early.
+                    # print(f'outvals: {outvals} and symbolic_outvals: {symbolic_outvals}')
                     assert numpy.allclose(
                         numpy.array(outvals), symbolic_outvals, equal_nan=True
                     )

neurallogic/symbolic_operator.py

@@ -8,11 +8,21 @@ def symbolic_operator(operator: str, x: str) -> str:
     return f'{operator}({x})'.replace('\'', '')
 
 
+@dispatch
+def symbolic_operator(operator: str, x: str, y: str):
+    return f'{operator}({x}, {y})'.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: int, 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))
@@ -28,11 +38,6 @@ 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)
@@ -86,4 +91,3 @@ def symbolic_operator(operator: str, x: numpy.ndarray):
 @dispatch
 def symbolic_operator(operator: str, x: list):
     return symbolic_operator(operator, numpy.array(x))
-

neurallogic/symbolic_primitives.py

@@ -2,6 +2,7 @@
 
 import jax
 import jax._src.lax_reference as lax_reference
+import jax._src.lax.lax as lax
 import numpy
 
 from neurallogic import symbolic_operator, symbolic_representation
@@ -59,31 +60,54 @@ def symbolic_gt(*args, **kwargs):
 
 
 def symbolic_abs(*args, **kwargs):
-    return symbolic(lax_reference.abs, 'numpy.absolute', *args, **kwargs)
+    return symbolic(lax_reference.abs, 'lax_reference.abs', *args, **kwargs)
+
+
+def symbolic_floor(*args, **kwargs):
+    return symbolic(lax_reference.floor, 'lax_reference.floor', *args, **kwargs)
+
+
+def symbolic_ceil(*args, **kwargs):
+    return symbolic(lax_reference.ceil, 'lax_reference.ceil', *args, **kwargs)
+
+
+def symbolic_round(*args, **kwargs):
+    # The reference implementation only supports away from zero
+    if kwargs['rounding_method'] == lax.RoundingMethod.AWAY_FROM_ZERO:
+        return symbolic(lax_reference.round, 'lax_reference.round', *args)
+    elif kwargs['rounding_method'] == lax.RoundingMethod.TO_NEAREST_EVEN:
+        return symbolic(numpy.around, 'numpy.around', *args)
+    else:
+        raise NotImplementedError(
+            f'rounding_method {str(kwargs["rounding_method"])} not implemented')
 
 
 def symbolic_add(*args, **kwargs):
-    return symbolic(lax_reference.add, 'numpy.add', *args, **kwargs)
+    return symbolic(lax_reference.add, 'lax_reference.add', *args, **kwargs)
 
 
 def symbolic_sub(*args, **kwargs):
-    return symbolic(lax_reference.sub, 'numpy.subtract', *args, **kwargs)
+    return symbolic(lax_reference.sub, 'lax_reference.sub', *args, **kwargs)
 
 
 def symbolic_mul(*args, **kwargs):
-    return symbolic(lax_reference.mul, 'numpy.multiply', *args, **kwargs)
+    return symbolic(lax_reference.mul, 'lax_reference.mul', *args, **kwargs)
 
 
 def symbolic_div(*args, **kwargs):
     return symbolic(lax_reference.div, 'lax_reference.div', *args, **kwargs)
 
 
+def symbolic_tan(*args, **kwargs):
+    return symbolic(lax_reference.tan, 'lax_reference.tan', *args, **kwargs)
+
+
 def symbolic_max(*args, **kwargs):
-    return symbolic(lax_reference.max, 'numpy.maximum', *args, **kwargs)
+    return symbolic(lax_reference.max, 'lax_reference.max', *args, **kwargs)
 
 
 def symbolic_min(*args, **kwargs):
-    return symbolic(lax_reference.min, 'numpy.minimum', *args, **kwargs)
+    return symbolic(lax_reference.min, 'lax_reference.min', *args, **kwargs)
 
 
 def symbolic_and(*args, **kwargs):
@@ -143,8 +167,11 @@ def symbolic_select_n(*args, **kwargs):
         # swap order of on_true and on_false
         return lax_reference.select(pred, on_false, on_true)
     else:
+        # TODO: to retain tensor structure we need to push down the select to the
+        # lowest level of the symbolic expression tree. This is not currently
+        # implemented.
+        print('WARNING: symbolic_select_n is not fully implemented. This may not work as expected.')
         # swap order of on_true and on_false
-        # TODO: need a more general solution to unquoting symbolic strings
         evaluable_pred = symbolic_representation.symbolic_representation(pred)
         evaluable_on_true = symbolic_representation.symbolic_representation(
             on_true)

neurallogic/symbolic_representation.py

@@ -1,6 +1,7 @@
 import numpy
 from plum import dispatch
 
+# TODO: need a more general solution to unquoting symbolic strings
 
 @dispatch
 def symbolic_representation(x: numpy.ndarray):

tests/test_mnist.py

@@ -47,8 +47,8 @@ def nln(type, x, width):
 def nln(type, x):
     num_classes = 10
 
-    x = hard_or.or_layer(type)(1000, nn.initializers.uniform(1.0), dtype=jax.numpy.float16)(x)
-    x = hard_not.not_layer(type)(num_classes)(x)
+    x = hard_or.or_layer(type)(1800, nn.initializers.uniform(1.0), dtype=jax.numpy.float16)(x)
+    x = hard_not.not_layer(type)(1, dtype=jax.numpy.float16)(x)
     x = x.ravel()
     x = harden_layer.harden_layer(type)(x) 
     x = x.reshape((num_classes, int(x.shape[0] / num_classes))) 
@@ -128,13 +128,14 @@ def get_datasets():
     ds_builder.download_and_prepare()
     train_ds = tfds.as_numpy(ds_builder.as_dataset(split="train", batch_size=-1))
     test_ds = tfds.as_numpy(ds_builder.as_dataset(split="test", batch_size=-1))
-    train_ds["image"] = jnp.float32(train_ds["image"]) / 255.0
-    test_ds["image"] = jnp.float32(test_ds["image"]) / 255.0
+    # XXXX
+    train_ds["image"] = (jnp.float32(train_ds["image"]) / 255.0)
+    test_ds["image"] = (jnp.float32(test_ds["image"]) / 255.0)
     # TODO: we don't need to do this even when we don't use the real encoder
     # Use grayscale information
     # Convert the floating point values in [0,1] to binary values in {0,1}
-    train_ds["image"] = jnp.round(train_ds["image"])
-    test_ds["image"] = jnp.round(test_ds["image"])
+    #train_ds["image"] = jnp.round(train_ds["image"])
+    #test_ds["image"] = jnp.round(test_ds["image"])
     return train_ds, test_ds