update saving and loading utils, making hebbian synapse use these utils for custom optimizer params saving and loading

Author: rxng8

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -11,6 +11,7 @@
 from ngclearn.components.synapses import DenseSynapse
 from ngclearn.utils import tensorstats
 from ngcsimlib import deprecate_args
+from ngclearn.utils.io_utils import save_pkl, load_pkl
 
 @partial(jit, static_argnums=[3, 4, 5, 6, 7, 8, 9])
 def _calc_update(
@@ -217,16 +218,12 @@ def __init__(
     def save(self, directory: str):
         super().save(directory)
         # Also save the optimizer parameters
-        file_name = directory + "/" + self.name + "_opt_params" + ".pkl"
-        with open(file_name, 'wb') as f:
-            pickle.dump(self.opt_params.get(), f)
+        save_pkl(directory, self.name + "_opt_params", self.opt_params.get())
 
     def load(self, directory: str):
         super().load(directory)
-        file_name = directory + "/" + self.name + "_opt_params" + ".pkl"
-        with open(file_name, 'rb') as f:
-            data = pickle.load(f)
-        self.opt_params.set(data)
+        # load the optimizer parameters in a custom way
+        self.opt_params.set(load_pkl(directory, self.name + "_opt_params"))
 
     @staticmethod
     def _compute_update(

ngclearn/utils/io_utils.py

@@ -4,6 +4,7 @@
 # import jax
 # from jax import numpy as jnp, grad, jit, vmap, random, lax
 import os, sys, pickle
+from typing import Any
 
 def serialize(fname, object): ## object "saving" routine
     """
@@ -65,3 +66,15 @@ def makedirs(directories):
     """
     for dir in directories:
         makedir(dir)
+
+
+def save_pkl(directory: str, name: str, value: Any) -> None:
+  file_name = directory + "/" + name + ".pkl"
+  with open(file_name, 'wb') as f:
+    pickle.dump(value, f)
+
+def load_pkl(directory: str, name: str) -> Any:
+  file_name = directory + "/" + name + ".pkl"
+  with open(file_name, 'rb') as f:
+    data = pickle.load(f)
+  return data