conditionally harden weights

Author: Z80coder

neurallogic/hard_and.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import numpy
 import jax
 from flax import linen as nn
 from typing import Callable
@@ -70,7 +69,7 @@ class SoftAndLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', self.weights_init,
+        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)
@@ -90,7 +89,7 @@ class HardAndLayer(nn.Module):
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
         weights = self.param(
-            'weights', nn.initializers.constant(0.0), weights_shape)
+            'bit_weights', nn.initializers.constant(0.0), weights_shape)
         return hard_and_layer(weights, x)
 
 

neurallogic/hard_not.py

@@ -41,7 +41,7 @@ class SoftNotLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param("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_not_layer(weights, x)
 
@@ -52,7 +52,7 @@ class HardNotLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param("weights", nn.initializers.constant(0.0), weights_shape)
+        weights = self.param("bit_weights", nn.initializers.constant(0.0), weights_shape)
         return hard_not_layer(weights, x)
 
 

neurallogic/hard_or.py

@@ -56,7 +56,7 @@ class SoftOrLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('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_or_layer(weights, x)
 
@@ -66,7 +66,7 @@ class HardOrLayer(nn.Module):
     @nn.compact
     def __call__(self, x):
         weights_shape = (self.layer_size, jax.numpy.shape(x)[-1])
-        weights = self.param('weights', nn.initializers.constant(0.0), weights_shape)
+        weights = self.param('bit_weights', nn.initializers.constant(0.0), weights_shape)
         return hard_or_layer(weights, x)
 
 class SymbolicOrLayer:

neurallogic/harden.py

@@ -36,14 +36,20 @@ def harden(x: jax.numpy.ndarray):
 
 @dispatch
 def harden(x: dict):
-    return symbolic_primitives.map_at_elements(x, harden)
+    # Only harden parameters that explicitly represent bits
+    def conditional_harden(k, v):
+        if k.startswith("bit_"):
+            return symbolic_primitives.map_at_elements(v, harden)
+        elif isinstance(v, dict) or isinstance(v, flax.core.FrozenDict) or isinstance(v, list):
+            return harden(v)
+        return v
+
+    return {k: conditional_harden(k, v) for k, v in x.items()}
 
 
 @dispatch
 def harden(x: flax.core.FrozenDict):
-    return flax.core.FrozenDict(
-        symbolic_primitives.map_at_elements(x.unfreeze(), harden)
-    )
+    return harden(x.unfreeze())
 
 
 @dispatch

neurallogic/real_encoder.py

@@ -66,6 +66,7 @@ def __call__(self, x):
         jaxpr = symbolic_generation.make_symbolic_flax_jaxpr(
             self.hard_real_encoder_layer, x
         )
+        print(f'SymbolicRealEncoderLayer: jaxpr:\n{jaxpr}')
         return symbolic_generation.symbolic_expression(jaxpr, x)
 
 

neurallogic/symbolic_generation.py

@@ -60,7 +60,7 @@ def put_variable(self, col: str, name: str, value: Any):
 
 
 def make_symbolic_flax_jaxpr(flax_layer, x):
-    actual_weights = flax_layer.get_variable("params", "weights")
+    actual_weights = flax_layer.get_variable("params", "bit_weights")
     # Convert actual weights to dummy numeric weights (if needed)
     if isinstance(actual_weights, list) or (
         isinstance(actual_weights, numpy.ndarray) and actual_weights.dtype == object
@@ -69,7 +69,7 @@ def make_symbolic_flax_jaxpr(flax_layer, x):
             actual_weights, lambda x: 0
         )
         numeric_weights = numpy.asarray(numeric_weights, dtype=numpy.int32)
-        put_variable(flax_layer, "params", "weights", numeric_weights)
+        put_variable(flax_layer, "params", "bit_weights", numeric_weights)
     # Convert input to dummy numeric input (if needed)
     if isinstance(x, list) or (isinstance(x, numpy.ndarray) and x.dtype == object):
         x = symbolic_primitives.map_at_elements(x, lambda x: 0)
@@ -171,9 +171,9 @@ def eval_jaxpr_impl(jaxpr):
                     symbolic_outvals = [symbolic_outvals]
                 if not symbolic:
                     # Check that the concrete and symbolic values are equal
-                    #print(
+                    # print(
                     #    f"outvals: {outvals} and symbolic_outvals: {symbolic_outvals}"
-                    #)
+                    # )
                     assert numpy.allclose(
                         numpy.array(outvals), symbolic_outvals, equal_nan=True
                     )