refactor some utils

Author: Z80coder

neurallogic/symbolic_generation.py

@@ -1,8 +1,6 @@
-import sys
 import typing
 from typing import Any, Mapping
 
-import flax
 import jax
 import numpy
 from jax import core
@@ -209,70 +207,6 @@ def eval_jaxpr_impl(jaxpr):
         return safe_map(symbolic_read, jaxpr.outvars)[0]
 
 
-def to_string(x):
-    return str(x)
-
-# TODO: use union types to consolidate these functions
-@dispatch
-def make_symbolic(x: dict):
-    return symbolic_primitives.map_at_elements(
-        x, to_string
-    )
-
-
-@dispatch
-def make_symbolic(x: list):
-    return symbolic_primitives.map_at_elements(
-        x, to_string
-    )
-
-
-@dispatch
-def make_symbolic(x: numpy.ndarray):
-    return symbolic_primitives.map_at_elements(
-        x, to_string
-    )
-
-
-@dispatch
-def make_symbolic(x: jax.numpy.ndarray):
-    return symbolic_primitives.map_at_elements(
-        convert_jax_to_numpy_arrays(x), to_string
-    )
-
-
-@dispatch
-def make_symbolic(x: bool):
-    return to_string(x)
-
-
-@dispatch
-def make_symbolic(x: str):
-    return to_string(x)
-
-
-@dispatch
-def convert_jax_to_numpy_arrays(x: jax.numpy.ndarray):
-    return numpy.asarray(x)
-
-
-@dispatch
-def convert_jax_to_numpy_arrays(x: dict):
-    return {k: convert_jax_to_numpy_arrays(v) for k, v in x.items()}
-
-
-@dispatch
-def make_symbolic(x: flax.core.FrozenDict):
-    x = convert_jax_to_numpy_arrays(x.unfreeze())
-    return flax.core.FrozenDict(make_symbolic(x))
-
-
-@dispatch
-def make_symbolic(*args):
-    return tuple([make_symbolic(arg) for arg in args])
-
-
-@dispatch
 def make_symbolic_jaxpr(func: typing.Callable, *args):
     return jax.make_jaxpr(lambda *args: func(*args))(*args)
 

tests/test_hard_and.py

@@ -1,9 +1,9 @@
 import jax
+import numpy
 import optax
 from flax import linen as nn
 from flax.training import train_state
 from jax import random
-import numpy
 
 from neurallogic import hard_and, harden, neural_logic_net, symbolic_generation
 from tests import utils
@@ -185,7 +185,7 @@ def test_net(type, x):
 
     # Compute symbolic result with symbolic inputs and symbolic weights, but where the symbols can be evaluated
     symbolic_input = ['True', 'False']
-    symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+    symbolic_weights = utils.make_symbolic(hard_weights)
     symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
     symbolic_output = symbolic_generation.eval_symbolic_expression(
         symbolic_output)

tests/test_hard_not.py

@@ -1,9 +1,9 @@
 import jax
 import jax.numpy as jnp
+import numpy
 import optax
 from flax.training import train_state
 from jax import random
-import numpy
 
 from neurallogic import hard_not, harden, neural_logic_net, symbolic_generation
 from tests import utils
@@ -250,7 +250,7 @@ def test_net(type, x):
 
     # Compute symbolic result with symbolic inputs and symbolic weights, but where the symbols can be evaluated
     symbolic_input = ["True", "False"]
-    symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+    symbolic_weights = utils.make_symbolic(hard_weights)
     symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
     symbolic_output = symbolic_generation.eval_symbolic_expression(
         symbolic_output)

tests/test_hard_or.py

@@ -1,10 +1,10 @@
 import jax
 import jax.numpy as jnp
+import numpy
 import optax
 from flax import linen as nn
 from flax.training import train_state
 from jax import random
-import numpy
 
 from neurallogic import hard_or, harden, neural_logic_net, symbolic_generation
 from tests import utils
@@ -184,7 +184,7 @@ def test_net(type, x):
 
     # Compute symbolic result with symbolic inputs and symbolic weights, but where the symbols can be evaluated
     symbolic_input = ['True', 'False']
-    symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+    symbolic_weights = utils.make_symbolic(hard_weights)
     symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
     symbolic_output = symbolic_generation.eval_symbolic_expression(
         symbolic_output)

tests/test_harden.py

@@ -1,6 +1,7 @@
+import flax
 import jax.numpy as jnp
+
 from neurallogic import harden
-import flax
 
 
 def test_harden_float():

tests/test_harden_layer.py

@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 
-from neurallogic import harden_layer, harden
+from neurallogic import harden, harden_layer
 
 
 def test_harden_layer():

tests/test_mnist.py

@@ -1,17 +1,18 @@
-from tqdm import tqdm
-from matplotlib import pyplot as plt
-import tensorflow as tf
-import tensorflow_datasets as tfds
 import jax
 import jax.numpy as jnp
+import ml_collections
 import numpy as np
+import optax
+import tensorflow as tf
+import tensorflow_datasets as tfds
 from flax import linen as nn
 from flax.metrics import tensorboard
 from flax.training import train_state
-import ml_collections
-from neurallogic import hard_not, hard_or, harden, harden_layer, neural_logic_net
-import optax
+from matplotlib import pyplot as plt
+from tqdm import tqdm
 
+from neurallogic import (hard_not, hard_or, harden, harden_layer,
+                         neural_logic_net)
 
 """
 MNIST test.

tests/test_network.py

@@ -9,6 +9,7 @@
 
 from neurallogic import (hard_and, hard_not, hard_or, harden, harden_layer,
                          neural_logic_net, symbolic_generation)
+from tests import utils
 
 config.update("jax_enable_x64", True)
 
@@ -70,7 +71,7 @@ def test_net(type, x):
         hard_result = hard.apply(hard_weights, hard_input)
         assert jnp.array_equal(hard_result, hard_expected)
 
-        symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+        symbolic_weights = utils.make_symbolic(hard_weights)
         symbolic_result = symbolic.apply(symbolic_weights, hard_input)
         symbolic_result = symbolic_generation.eval_symbolic_expression(
             symbolic_result)

tests/test_real_encoder.py

@@ -9,6 +9,7 @@
 
 from neurallogic import (harden, neural_logic_net, real_encoder,
                          symbolic_generation)
+from tests import utils
 
 # Uncomment to debug NaNs
 # config.update("jax_debug_nans", True)
@@ -249,7 +250,7 @@ def test_net(type, x):
 
     # Compute symbolic result with symbolic inputs and symbolic weights, but where the symbols can be evaluated
     symbolic_input = ['1.0', '0.0']
-    symbolic_weights = symbolic_generation.make_symbolic(hard_weights)
+    symbolic_weights = utils.make_symbolic(hard_weights)
     symbolic_output = symbolic.apply(symbolic_weights, symbolic_input)
     symbolic_output = symbolic_generation.eval_symbolic_expression(
         symbolic_output)

tests/test_symbolic_generation.py

@@ -1,8 +1,11 @@
-from neurallogic import neural_logic_net, harden, harden_layer, hard_or, hard_not, symbolic_generation, symbolic_primitives
-from tests import test_mnist
-import numpy
 import jax
 import jax.numpy as jnp
+import numpy
+
+from neurallogic import (hard_not, hard_or, harden, harden_layer,
+                         neural_logic_net, symbolic_generation,
+                         symbolic_primitives)
+from tests import test_mnist, utils
 
 
 def nln(type, x, width):
@@ -45,7 +48,7 @@ def test_symbolic_generation():
     assert numpy.array_equal(symbolic_output, hard_output)
 
     # Check the standard evaluation of the network equals the non-standard symbolic evaluation
-    symbolic_mock_input = symbolic_generation.make_symbolic(hard_mock_input)
+    symbolic_mock_input = utils.make_symbolic(hard_mock_input)
     symbolic_output = symbolic.apply(hard_weights, symbolic_mock_input)
     assert numpy.array_equal(hard_output.shape, symbolic_output.shape)