document problem with 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_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
 
 

tests/test_real_encoder.py

@@ -4,13 +4,16 @@
 import numpy
 
 from neurallogic import (hard_and, hard_majority, hard_not, hard_or, hard_xor,
-                         harden, harden_layer, neural_logic_net,
+                         harden, harden_layer, neural_logic_net, real_encoder,
                          symbolic_generation)
 from tests import utils
 
 
 def nln(type, x, width):
-    # TODO: test real_encoder layer
+    # Can't symbolically support this layer yet since the symbolic output is an unevaluated string that
+    # lacks the correct tensor structure
+    # x = real_encoder.real_encoder_layer(type)(2)(x)
+    # x = x.ravel()
     x = hard_or.or_layer(type)(width)(x)
     x = hard_and.and_layer(type)(width)(x)
     x = hard_xor.xor_layer(type)(width)(x)