test fixes and black

elegy/hooks_test.py

@@ -32,7 +32,7 @@ def test_summaries(self):
             elegy.hooks.add_summary(("a", 0, "b"), None, 2.0)
             summaries = elegy.hooks.get_summaries()
 
-        assert summaries[0] == (("a", 0, "b"), None, 2.0)
+        assert summaries[0] == (("a", 0, "b"), None, 2.0, None)
 
     def test_no_summaries(self):
         assert not elegy.hooks.summaries_active()
@@ -65,4 +65,4 @@ def f(x):
         assert x == 6
         assert losses["x_loss"] == 6
         assert metrics["x"] == 7
-        assert summaries[0] == (("a", 0, "b"), jax.nn.relu, 8)
+        assert summaries[0] == (("a", 0, "b"), jax.nn.relu, 8, None)

elegy/module.py

@@ -382,7 +382,13 @@ def __call__(self, *args, **kwargs) -> tp.Any:
     def call(self, *args, **kwargs):
         ...
 
-    def add_summary(self, name: str, f: tp.Any, value: tp.Any, input_values: tp.Optional[tp.Tuple[tp.Tuple, tp.Dict]] = None):
+    def add_summary(
+        self,
+        name: str,
+        f: tp.Any,
+        value: tp.Any,
+        input_values: tp.Optional[tp.Tuple[tp.Tuple, tp.Dict]] = None,
+    ):
         if hooks.summaries_active():
             path = get_module_path(self) + (name,)
             assert path is not None
@@ -623,7 +629,6 @@ def slice(
             self, start_module, end_module, sample_input
         )
 
-
     def update_parameter(self, name: str, value: tp.Any) -> None:
         """
         Update a parameter of the current module.

elegy/module_slicing.py

@@ -52,7 +52,7 @@ class Edge:
     def __init__(self, summary: elegy.types.SummaryTableEntry):
         self.module = summary.module
         # standardize paths with a leading '/'
-        self.modulename = '/'+summary.path
+        self.modulename = "/" + summary.path
         # convert the output and input arrays in the summary to unique IDs as returned by id()
         self.output_ids = jax.tree_leaves(jax.tree_map(id, summary.output_value))
         self.input_ids = jax.tree_map(id, summary.input_value)

elegy/module_slicing_test.py

@@ -25,7 +25,7 @@ def test_basic_slice_by_name(self):
             submodule = basicmodule.slice(start, end, x)
             submodel = elegy.Model(submodule)
             basicmodule.init(rng=elegy.RNGSeq(0), set_defaults=True)(x)
-            #submodel.summary(x)
+            # submodel.summary(x)
             assert submodel.predict(x).shape == (32, 10)
             assert jnp.all(submodel.predict(x) == basicmodule.test_call(x))
 

elegy/module_test.py

@@ -214,11 +214,13 @@ def call(self, x):
                 ("ais/1",),
                 m.ais[1],
                 12,
+                ((2.0,), {}),
             ),
             (
                 (),
                 m,
                 13,
+                ((2.0,), {}),
             ),
         ]
         assert parameters == {
@@ -256,11 +258,13 @@ def call(self, x):
                 ("a_1",),
                 m.a_1,
                 12,
+                ((2.0,), {}),
             ),
             (
                 (),
                 m,
                 13,
+                ((2.0,), {}),
             ),
         ]
         assert params == {

elegy/types.py

@@ -121,11 +121,11 @@ class Summary(tp.NamedTuple):
     input_values: tp.Union[tp.Tuple[tp.Tuple, tp.Dict], None] = None
 
     def tree_flatten(self):
-        return ((self.value,self.input_values), (self.path, self.module))
+        return ((self.value, self.input_values), (self.path, self.module))
 
     @classmethod
     def tree_unflatten(cls, aux_data, children):
-        (value,input_values) = children
+        (value, input_values) = children
         path, module = aux_data
 
         return cls(path, module, value, input_values)
@@ -168,7 +168,7 @@ def totals_entry(
 
     def tree_flatten(self):
         return (
-            (self.output_value,self.input_value),
+            (self.output_value, self.input_value),
             (
                 self.path,
                 self.module_type_name,
@@ -191,7 +191,7 @@ def tree_unflatten(cls, aux_data, children):
             non_trainable_params_count,
             non_trainable_params_size,
         ) = aux_data
-        (output_value,input_value) = children
+        (output_value, input_value) = children
 
         return cls(
             path=path,