maintain test_harden

Author: Z80coder

neurallogic/harden.py

@@ -12,6 +12,10 @@ def harden_float(x: float) -> bool:
 def harden(x: float):
     return harden_float(x)
 
+@dispatch
+def harden(x: list):
+    return symbolic_primitives.map_at_elements(x, harden_float)
+
 @dispatch
 def harden(x: numpy.ndarray):
     return symbolic_primitives.map_at_elements(x, harden_float)
@@ -24,6 +28,10 @@ def harden(x: jax.numpy.ndarray):
 def harden(x: dict):
     return symbolic_primitives.map_at_elements(x, harden_float)
 
+@dispatch
+def harden(x: flax.core.FrozenDict):
+    return flax.core.FrozenDict(symbolic_primitives.map_at_elements(x.unfreeze(), harden_float))
+
 @dispatch
 def harden(*args):
     if len(args) == 1:
@@ -34,9 +42,5 @@ def harden(*args):
 def map_keys_nested(f, d: dict) -> dict:
     return {f(k): map_keys_nested(f, v) if isinstance(v, dict) else v for k, v in d.items()}
 
-
 def hard_weights(weights):
-    unfrozen_weights = weights.unfreeze()
-    hard_weights = harden(unfrozen_weights)
-    return flax.core.FrozenDict(map_keys_nested(lambda str: str.replace("Soft", "Hard"), hard_weights))
-
+    return flax.core.FrozenDict(map_keys_nested(lambda str: str.replace("Soft", "Hard"), harden(weights.unfreeze())))

tests/test_harden.py

@@ -12,26 +12,26 @@ def test_harden_float():
 
 
 def test_harden_list():
-    assert harden.harden_list([0.5, 0.6, 0.4, 0.0, 1.0]) == [
+    assert harden.harden([0.5, 0.6, 0.4, 0.0, 1.0]) == [
         False, True, False, False, True]
 
 
 def test_harden_array():
-    assert jnp.array_equal(harden.harden_array(
+    assert jnp.array_equal(harden.harden(
         jnp.array([0.5, 0.6, 0.4, 0.0, 1.0])), [False, True, False, False, True])
 
 
 def test_harden_dict():
     dict = {'a': 0.5, 'b': 0.6, 'c': 0.4, 'd': 0.0, 'e': 1.0}
     expected_dict = {'a': False, 'b': True, 'c': False, 'd': False, 'e': True}
-    assert harden.harden_dict(dict) == expected_dict
+    assert harden.harden(dict) == expected_dict
 
 
 def test_harden_frozen_dict():
     dict = flax.core.frozen_dict.FrozenDict(
         {'a': 0.5, 'b': 0.6, 'c': 0.4, 'd': 0.0, 'e': 1.0})
     expected_dict = {'a': False, 'b': True, 'c': False, 'd': False, 'e': True}
-    assert harden.harden_dict(dict) == expected_dict
+    assert harden.harden(dict) == expected_dict
 
 
 def test_harden():