Merge pull request #57 from github/noisy-xor

Infrastructure for experiments

Author: Z80coder

neurallogic/hard_and.py

@@ -1,38 +1,19 @@
-from typing import Any
+from typing import Callable
 
 import jax
 from flax import linen as nn
-from typing import Callable
-
-
-from neurallogic import neural_logic_net, symbolic_generation
-
-
-def soft_and_include(w: float, x: float) -> float:
-    """
-    w > 0.5 implies the and operation is active, else inactive
-
-    Assumes x is in [0, 1]
-
-    Corresponding hard logic: x OR ! w
-    """
-    w = jax.numpy.clip(w, 0.0, 1.0)
-    return jax.numpy.maximum(x, 1.0 - w)
-
-
-
-def hard_and_include(w, x):
-    return jax.numpy.logical_or(x, jax.numpy.logical_not(w))
 
+from neurallogic import hard_masks, neural_logic_net, symbolic_generation
 
 
+# TODO: seperate and operation from mask operation
 def soft_and_neuron(w, x):
-    x = jax.vmap(soft_and_include, 0, 0)(w, x)
+    x = jax.vmap(hard_masks.soft_mask_to_true, 0, 0)(w, x)
     return jax.numpy.min(x)
 
 
 def hard_and_neuron(w, x):
-    x = jax.vmap(hard_and_include, 0, 0)(w, x)
+    x = jax.vmap(hard_masks.hard_mask_to_true, 0, 0)(w, x)
     return jax.lax.reduce(x, True, jax.lax.bitwise_and, [0])
 
 
@@ -41,6 +22,7 @@ def hard_and_neuron(w, x):
 hard_and_layer = jax.vmap(hard_and_neuron, (0, None), 0)
 
 
+# TODO: move initialization to separate file
 def initialize_near_to_zero():
     # TODO: investigate better initialization
     def init(key, shape, dtype):
@@ -51,6 +33,7 @@ def init(key, shape, dtype):
         x = 0.5 * x - 1
         x = jax.numpy.clip(x, 0.001, 0.999)
         return x
+
     return init
 
 
@@ -62,15 +45,17 @@ class SoftAndLayer(nn.Module):
         layer_size: The number of neurons in the layer.
         weights_init: The initializer function for the weight matrix.
     """
+
     layer_size: int
     weights_init: Callable = initialize_near_to_zero()
     dtype: jax.numpy.dtype = jax.numpy.float32
 
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('bit_weights', self.weights_init,
-                             weights_shape, self.dtype)
+        weights = self.param(
+            "bit_weights", self.weights_init, weights_shape, self.dtype
+        )
         x = jax.numpy.asarray(x, self.dtype)
         return soft_and_layer(weights, x)
 
@@ -83,13 +68,15 @@ class HardAndLayer(nn.Module):
     Attributes:
         layer_size: The number of neurons in the layer.
     """
+
     layer_size: int
 
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'bit_weights', nn.initializers.constant(True), weights_shape)
+            "bit_weights", nn.initializers.constant(True), weights_shape
+        )
         return hard_and_layer(weights, x)
 
 
@@ -104,6 +91,13 @@ def __call__(self, x):
 
 
 and_layer = neural_logic_net.select(
-    lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: SoftAndLayer(layer_size, weights_init, dtype),
-    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: HardAndLayer(layer_size),
-    lambda layer_size, weights_init=nn.initializers.constant(True), dtype=jax.numpy.float32: SymbolicAndLayer(layer_size))
+    lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: SoftAndLayer(
+        layer_size, weights_init, dtype
+    ),
+    lambda layer_size, weights_init=nn.initializers.constant(
+        True
+    ), dtype=jax.numpy.float32: HardAndLayer(layer_size),
+    lambda layer_size, weights_init=nn.initializers.constant(
+        True
+    ), dtype=jax.numpy.float32: SymbolicAndLayer(layer_size),
+)

neurallogic/hard_dropout.py

@@ -0,0 +1,90 @@
+from typing import Optional, Sequence
+
+import jax
+from flax import linen as nn
+from jax import lax, random
+
+from neurallogic import neural_logic_net
+
+
+class SoftHardDropout(nn.Module):
+    """Create a dropout layer suitable for dropping soft-bit values.
+    Adapted from flax/stochastic.py
+
+
+    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
+    to include an RNG seed named `'dropout'`. For example::
+
+      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
+
+    Attributes:
+      rate: the dropout probability.  (_not_ the keep rate!)
+      broadcast_dims: dimensions that will share the same dropout mask
+      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
+        masked, whereas if true, no mask is applied and the inputs are returned
+        as is.
+      rng_collection: the rng collection name to use when requesting an rng key.
+    """
+
+    rate: float
+    broadcast_dims: Sequence[int] = ()
+    deterministic: Optional[bool] = None
+    rng_collection: str = "dropout"
+    dropout_value: float = 0.0
+    dtype: jax.numpy.dtype = jax.numpy.float32
+
+    @nn.compact
+    def __call__(self, inputs, deterministic: Optional[bool] = None):
+        """Applies a random dropout mask to the input.
+
+        Args:
+          inputs: the inputs that should be randomly masked.
+          Masking means setting the input bits to 0.5.
+          deterministic: if false the inputs are masked,
+          whereas if true, no mask is applied and the inputs are returned
+          as is.
+
+        Returns:
+          The masked inputs
+        """
+        deterministic = nn.merge_param(
+            "deterministic", self.deterministic, deterministic
+        )
+
+        if (self.rate == 0.0) or deterministic:
+            return inputs
+
+        # Prevent gradient NaNs in 1.0 edge-case.
+        if self.rate == 1.0:
+            return jax.numpy.zeros_like(inputs)
+
+        keep_prob = 1.0 - self.rate
+        rng = self.make_rng(self.rng_collection)
+        broadcast_shape = list(inputs.shape)
+        for dim in self.broadcast_dims:
+            broadcast_shape[dim] = 1
+        mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
+        mask = jax.numpy.broadcast_to(mask, inputs.shape)
+        masked_values = jax.numpy.full_like(
+            inputs, self.dropout_value, dtype=self.dtype
+        )
+        return lax.select(mask, inputs, masked_values)
+
+
+class HardHardDropout(nn.Module):
+    @nn.compact
+    def __call__(self, inputs, deterministic: Optional[bool] = None):
+        return inputs
+
+
+class SymbolicHardDropout(nn.Module):
+    @nn.compact
+    def __call__(self, inputs, deterministic: Optional[bool] = None):
+        return inputs
+
+
+hard_dropout = neural_logic_net.select(
+    lambda **kwargs: SoftHardDropout(**kwargs),
+    lambda **kwargs: HardHardDropout(**kwargs),
+    lambda **kwargs: SymbolicHardDropout(**kwargs),
+)

neurallogic/hard_masks.py

@@ -0,0 +1,129 @@
+from typing import Callable
+
+import jax
+from flax import linen as nn
+
+from neurallogic import neural_logic_net, symbolic_generation
+
+
+def soft_mask_to_true(w: float, x: float) -> float:
+    """
+    w > 0.5 implies the mask operation is inactive, else active
+
+    Assumes x is in [0, 1]
+
+    Corresponding hard logic: x OR ! w
+    """
+    w = jax.numpy.clip(w, 0.0, 1.0)
+    return jax.numpy.maximum(x, 1.0 - w)
+
+
+def hard_mask_to_true(w, x):
+    return jax.numpy.logical_or(x, jax.numpy.logical_not(w))
+
+
+soft_mask_to_true_neuron = jax.vmap(soft_mask_to_true, 0, 0)
+
+hard_mask_to_true_neuron = jax.vmap(hard_mask_to_true, 0, 0)
+
+
+soft_mask_to_true_layer = jax.vmap(soft_mask_to_true_neuron, (0, None), 0)
+
+hard_mask_to_true_layer = jax.vmap(hard_mask_to_true_neuron, (0, None), 0)
+
+
+def soft_mask_to_false(w: float, x: float) -> float:
+    """
+    w > 0.5 implies the mask is inactive, else active
+
+    Assumes x is in [0, 1]
+
+    Corresponding hard logic: b AND w
+    """
+    w = jax.numpy.clip(w, 0.0, 1.0)
+    return 1.0 - jax.numpy.maximum(1.0 - x, 1.0 - w)
+
+
+def hard_mask_to_false(w, x):
+    return jax.numpy.logical_and(x, w)
+
+
+soft_mask_to_false_neuron = jax.vmap(soft_mask_to_false, 0, 0)
+
+hard_mask_to_false_neuron = jax.vmap(hard_mask_to_false, 0, 0)
+
+
+soft_mask_to_false_layer = jax.vmap(soft_mask_to_false_neuron, (0, None), 0)
+
+hard_mask_to_false_layer = jax.vmap(hard_mask_to_false_neuron, (0, None), 0)
+
+
+class SoftMaskLayer(nn.Module):
+    mask_layer_operation: Callable
+    layer_size: int
+    weights_init: Callable = nn.initializers.uniform(1.0)
+    dtype: jax.numpy.dtype = jax.numpy.float32
+
+    @nn.compact
+    def __call__(self, x):
+        weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
+        weights = self.param(
+            "bit_weights", self.weights_init, weights_shape, self.dtype
+        )
+        x = jax.numpy.asarray(x, self.dtype)
+        return self.mask_layer_operation(weights, x)
+
+
+class HardMaskLayer(nn.Module):
+    mask_layer_operation: Callable
+    layer_size: int
+    weights_init: Callable = nn.initializers.constant(True)
+
+    @nn.compact
+    def __call__(self, x):
+        weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
+        weights = self.param("bit_weights", self.weights_init, weights_shape)
+        return self.mask_layer_operation(weights, x)
+
+
+class SymbolicMaskLayer:
+    def __init__(self, mask_layer):
+        self.hard_mask_layer = mask_layer
+
+    def __call__(self, x):
+        jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(self.hard_mask_layer, x)
+        return symbolic_generation.symbolic_expression(jaxpr, x)
+
+
+mask_to_true_layer = neural_logic_net.select(
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SoftMaskLayer(
+        soft_mask_to_true_layer, layer_size, weights_init, dtype
+    ),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: HardMaskLayer(hard_mask_to_true_layer, layer_size),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SymbolicMaskLayer(
+        HardMaskLayer(hard_mask_to_true_layer, layer_size)
+    ),
+)
+
+
+mask_to_false_layer = neural_logic_net.select(
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SoftMaskLayer(
+        soft_mask_to_false_layer, layer_size, weights_init, dtype
+    ),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: HardMaskLayer(hard_mask_to_false_layer, layer_size),
+    lambda layer_size, weights_init=nn.initializers.uniform(
+        1.0
+    ), dtype=jax.numpy.float32: SymbolicMaskLayer(
+        HardMaskLayer(hard_mask_to_false_layer, layer_size)
+    ),
+)