remove jaxpr layer type

Author: Z80coder

neurallogic/hard_and.py

@@ -6,7 +6,7 @@
 from typing import Callable
 
 
-from neurallogic import neural_logic_net, sym_gen, symbolic_primitives
+from neurallogic import neural_logic_net, sym_gen
 
 
 def soft_and_include(w: float, x: float) -> float:
@@ -42,7 +42,6 @@ def hard_and_neuron(w, x):
 hard_and_layer = jax.vmap(hard_and_neuron, (0, None), 0)
 
 
-
 def initialize_near_to_zero():
     # TODO: investigate better initialization
     def init(key, shape, dtype):
@@ -95,44 +94,17 @@ def __call__(self, x):
         return hard_and_layer(weights, x)
 
 
-class JaxprAndLayer:
-    def __init__(self, layer_size):
-        self.layer_size = layer_size
-        self.hard_and_layer = HardAndLayer(self.layer_size)
-
-    def __call__(self, x):
-        jaxpr = sym_gen.make_symbolic_jaxpr(self.hard_and_layer, x)
-        return sym_gen.eval_symbolic(jaxpr, x)
-
-
 class SymbolicAndLayer:
     def __init__(self, layer_size):
         self.layer_size = layer_size
         self.hard_and_layer = HardAndLayer(self.layer_size)
 
     def __call__(self, x):
-        actual_weights = self.hard_and_layer.get_variable("params", "weights")
-        if isinstance(actual_weights, list) or (isinstance(actual_weights, numpy.ndarray) and actual_weights.dtype == object):
-            numeric_weights = symbolic_primitives.map_at_elements(actual_weights, lambda x: 0)
-            numeric_weights = numpy.asarray(numeric_weights, dtype=numpy.float32)
-            sym_gen.put_variable(self.hard_and_layer, "params", "weights", numeric_weights)
-        if isinstance(x, list) or (isinstance(x, numpy.ndarray) and x.dtype == object):
-            xn = symbolic_primitives.map_at_elements(x, lambda x: 0)
-            xn = numpy.asarray(xn, dtype=numpy.float32)
-        else:
-            xn = x
-        jaxpr = sym_gen.make_symbolic_jaxpr(self.hard_and_layer, xn)
-        # Swap out the numeric consts (that represent the weights) for the symbolic weights
-        jaxpr.consts = [actual_weights]
+        jaxpr = sym_gen.make_symbolic_flax_jaxpr(self.hard_and_layer, x)
         return sym_gen.symbolic_expression(jaxpr, 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=initialize_near_to_zero(
-    ), dtype=jax.numpy.float32: HardAndLayer(layer_size),
-    lambda layer_size, weights_init=initialize_near_to_zero(
-    ), dtype=jax.numpy.float32: JaxprAndLayer(layer_size),
-    lambda layer_size, weights_init=initialize_near_to_zero(),
-    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=initialize_near_to_zero(), dtype=jax.numpy.float32: HardAndLayer(layer_size),
+    lambda layer_size, weights_init=initialize_near_to_zero(), dtype=jax.numpy.float32: SymbolicAndLayer(layer_size))

neurallogic/neural_logic_net.py

@@ -1,14 +1,13 @@
 from enum import Enum
 from flax import linen as nn
 
-NetType = Enum('NetType', ['Soft', 'Hard', 'Jaxpr', 'Symbolic'])
+NetType = Enum('NetType', ['Soft', 'Hard', 'Symbolic'])
 
-def select(soft, hard, jaxpr, symbolic):
+def select(soft, hard, symbolic):
     def selector(type: NetType):
         return {
             NetType.Soft: soft,
             NetType.Hard: hard,
-            NetType.Jaxpr: jaxpr,
             NetType.Symbolic: symbolic
         }[type]
     return selector
@@ -22,14 +21,9 @@ class HardNet(nn.Module):
         @nn.compact
         def __call__(self, x):
             return f(NetType.Hard, x)
-    class JaxprNet(nn.Module):
-        @nn.compact
-        def __call__(self, x):
-            return f(NetType.Jaxpr, x)
     class SymbolicNet(nn.Module):
         @nn.compact
         def __call__(self, x):
             return f(NetType.Symbolic, x)
-    return SoftNet(), HardNet(), JaxprNet(), SymbolicNet()
+    return SoftNet(), HardNet(), SymbolicNet()
 
-# TODO: support init of all three net types at once

neurallogic/sym_gen.py

@@ -50,6 +50,24 @@ def put_variable(self, col: str, name: str, value: Any):
     self.scope._variables = self.scope.variables().unfreeze()
     scope_put_variable(self.scope, col, name, value)
 
+def make_symbolic_flax_jaxpr(flax_layer, x):
+    actual_weights = flax_layer.get_variable("params", "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):
+        numeric_weights = symbolic_primitives.map_at_elements(actual_weights, lambda x: 0)
+        numeric_weights = numpy.asarray(numeric_weights, dtype=numpy.float32)
+        put_variable(flax_layer, "params", "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)
+        x = numpy.asarray(x, dtype=numpy.float32)
+    # Make the jaxpr that corresponds to the flax layer
+    jaxpr = make_symbolic_jaxpr(flax_layer, x)
+    # Replace the dummy numeric weights with the actual weights in the jaxpr
+    jaxpr.consts = [actual_weights]
+    return jaxpr
+
+
 def eval_jaxpr(symbolic, jaxpr, consts, *args):
     """Evaluates a jaxpr by interpreting it as Python code.
 

tests/test_hard_and.py

@@ -18,7 +18,7 @@ def check_consistency(soft: typing.Callable, hard: typing.Callable, expected, *a
     hard_expected = harden.harden(expected)
     assert numpy.allclose(hard(*hard_args), hard_expected)
 
-    # Check that the jaxpr expression performs as expected
+    # Check that the jaxpr performs as expected
     symbolic_f = sym_gen.make_symbolic_jaxpr(hard, *hard_args)
     assert numpy.allclose(sym_gen.eval_symbolic(
         symbolic_f, *hard_args), hard_expected)
@@ -91,7 +91,7 @@ def test_net(type, x):
         x = x.ravel()
         return x
 
-    soft, hard, jaxpr, symbolic = neural_logic_net.net(test_net)
+    soft, hard, symbolic = neural_logic_net.net(test_net)
     weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
     hard_weights = harden.hard_weights(weights)
 
@@ -124,10 +124,6 @@ def test_net(type, x):
         hard_output = hard.apply(hard_weights, hard_input)
         assert jax.numpy.allclose(hard_output, hard_expected)
 
-        # Check that the jaxpr expression performs as expected
-        jaxpr_output = jaxpr.apply(hard_weights, hard_input)
-        assert numpy.allclose(jaxpr_output, hard_expected)
-
         # Check that the symbolic function performs as expected
         symbolic_output = symbolic.apply(hard_weights, hard_input)
         assert numpy.allclose(symbolic_output, hard_expected)
@@ -137,7 +133,7 @@ def test_train_and():
     def test_net(type, x):
         return hard_and.and_layer(type)(4, nn.initializers.uniform(1.0))(x)
 
-    soft, hard, jaxpr, symbolic = neural_logic_net.net(test_net)
+    soft, hard, symbolic = neural_logic_net.net(test_net)
     soft_weights = soft.init(random.PRNGKey(0), [0.0, 0.0])
 
     x = [
@@ -184,7 +180,7 @@ def test_net(type, x):
         x = hard_and.and_layer(type)(4, nn.initializers.uniform(1.0))(x)
         return x
 
-    soft, hard, jaxpr, symbolic = neural_logic_net.net(test_net)
+    soft, hard, symbolic = neural_logic_net.net(test_net)
 
     # Compute soft result
     soft_input = jax.numpy.array([0.6, 0.45])