jinns 0.9.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

jinns/utils/_pinn.py

@@ -2,44 +2,53 @@
 Implements utility function to create PINNs
 """
 
-from typing import Callable
+from typing import Callable, Literal
+from dataclasses import InitVar
 import jax
 import jax.numpy as jnp
-from jax.typing import ArrayLike
 import equinox as eqx
 
+from jaxtyping import Array, Key, PyTree, Float
+
+from jinns.parameters._params import Params
+
 
 class _MLP(eqx.Module):
     """
     Class to construct an equinox module from a key and a eqx_list. To be used
-    in pair with the function `create_PINN`
+    in pair with the function `create_PINN`.
+
+    Parameters
+    ----------
+    key : InitVar[Key]
+        A jax random key for the layer initializations.
+    eqx_list : InitVar[tuple[tuple[Callable, int, int] | Callable, ...]]
+        A tuple of tuples of successive equinox modules and activation functions to
+        describe the PINN architecture. The inner tuples must have the eqx module or
+        activation function as first item, other items represents arguments
+        that could be required (eg. the size of the layer).
+        The `key` argument need not be given.
+        Thus typical example is `eqx_list=
+        ((eqx.nn.Linear, 2, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 20),
+            jax.nn.tanh,
+            (eqx.nn.Linear, 20, 1)
+        )`.
     """
 
-    layers: list
+    key: InitVar[Key] = eqx.field(kw_only=True)
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
+        kw_only=True
+    )
 
-    def __init__(self, key, eqx_list):
-        """
-        Parameters
-        ----------
-        key
-            A jax random key
-        eqx_list
-            A list of list of successive equinox modules and activation functions to
-            describe the PINN architecture. The inner lists have the eqx module or
-            axtivation function as first item, other items represents arguments
-            that could be required (eg. the size of the layer).
-            __Note:__ the `key` argument need not be given.
-            Thus typical example is `eqx_list=
-            [[eqx.nn.Linear, 2, 20],
-                [jax.nn.tanh],
-                [eqx.nn.Linear, 20, 20],
-                [jax.nn.tanh],
-                [eqx.nn.Linear, 20, 20],
-                [jax.nn.tanh],
-                [eqx.nn.Linear, 20, 1]
-            ]`
-        """
+    # NOTE that the following should NOT be declared as static otherwise the
+    # eqx.partition that we use in the PINN module will misbehave
+    layers: list[eqx.Module] = eqx.field(init=False)
 
+    def __post_init__(self, key, eqx_list):
         self.layers = []
         for l in eqx_list:
             if len(l) == 1:
@@ -48,75 +57,118 @@ class _MLP(eqx.Module):
                 key, subkey = jax.random.split(key, 2)
                 self.layers.append(l[0](*l[1:], key=subkey))
 
-    def __call__(self, t):
+    def __call__(self, t: Float[Array, "input_dim"]) -> Float[Array, "output_dim"]:
         for layer in self.layers:
             t = layer(t)
         return t
 
 
 class PINN(eqx.Module):
-    """
-    Basically a wrapper around the `__call__` function to be able to give a type to
-    our former `self.u`
-    The function create_PINN has the role to population the `__call__` function
-    """
+    r"""
+    A PINN object, i.e., a neural network compatible with the rest of jinns.
+    This is typically created with `create_PINN` which creates iternally a
+    `_MLP` object. However, a user could directly creates their PINN using this
+    class by passing a eqx.Module (for argument `mlp`)
+    that plays the role of the NN and that is
+    already instanciated.
 
-    slice_solution: ArrayLike
-    eq_type: str = eqx.field(static=True)
-    input_transform: Callable = eqx.field(static=True)
-    output_transform: Callable = eqx.field(static=True)
-    output_slice: ArrayLike
-    params: eqx.Module
-    static: eqx.Module
+    Parameters
+    ----------
+    slice_solution : slice
+        A jnp.s\_ object which indicates which axis of the PINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PINN output. This argument is useful
+        when the PINN is also used to output equation parameters for example
+        Note that it must be a slice and not an integer (a preprocessing of the
+        user provided argument takes care of it).
+    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
+        A string with three possibilities.
+        "ODE": the PINN is called with one input `t`.
+        "statio_PDE": the PINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
+        can be high dimensional.
+        **Note**: the input dimension as given in eqx_list has to match the sum
+        of the dimension of `t` + the dimension of `x` or the output dimension
+        after the `input_transform` function.
+    input_transform : Callable[[Float[Array, "input_dim"], Params], Float[Array, "output_dim"]]
+        A function that will be called before entering the PINN. Its output(s)
+        must match the PINN inputs (except for the parameters).
+        Its inputs are the PINN inputs (`t` and/or `x` concatenated together)
+        and the parameters. Default is no operation.
+    output_transform : Callable[[Float[Array, "input_dim"], Float[Array, "output_dim"], Params], Float[Array, "output_dim"]]
+        A function with arguments begin the same input as the PINN, the PINN
+        output and the parameter. This function will be called after exiting the PINN.
+        Default is no operation.
+    output_slice : slice, default=None
+        A jnp.s\_[] to determine the different dimension for the PINN.
+        See `shared_pinn_outputs` argument of `create_PINN`.
+    mlp : eqx.Module
+        The actual neural network instanciated as an eqx.Module.
+    """
 
-    def __init__(
-        self,
-        mlp,
-        slice_solution,
-        eq_type,
-        input_transform,
-        output_transform,
-        output_slice,
-    ):
+    slice_solution: slice = eqx.field(static=True, kw_only=True)
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+        static=True, kw_only=True
+    )
+    input_transform: Callable[
+        [Float[Array, "input_dim"], Params], Float[Array, "output_dim"]
+    ] = eqx.field(static=True, kw_only=True)
+    output_transform: Callable[
+        [Float[Array, "input_dim"], Float[Array, "output_dim"], Params],
+        Float[Array, "output_dim"],
+    ] = eqx.field(static=True, kw_only=True)
+    output_slice: slice = eqx.field(static=True, kw_only=True, default=None)
+
+    mlp: InitVar[eqx.Module] = eqx.field(kw_only=True)
+
+    params: PyTree = eqx.field(init=False)
+    static: PyTree = eqx.field(init=False, static=True)
+
+    def __post_init__(self, mlp):
         self.params, self.static = eqx.partition(mlp, eqx.is_inexact_array)
-        self.slice_solution = slice_solution
-        self.eq_type = eq_type
-        self.input_transform = input_transform
-        self.output_transform = output_transform
-        self.output_slice = output_slice
 
-    def init_params(self):
+    def init_params(self) -> PyTree:
+        """
+        Returns an initial set of parameters
+        """
         return self.params
 
-    def __call__(self, *args):
+    def __call__(self, *args) -> Float[Array, "output_dim"]:
+        """
+        Calls `eval_nn` with rearranged arguments
+        """
         if self.eq_type == "ODE":
             (t, params) = args
             if len(t.shape) == 0:
                 t = t[..., None]  #  Add mandatory dimension which can be lacking
                 # (eg. for the ODE batches) but this dimension can already
                 # exists (eg. for user provided observation times)
-            return self._eval_nn(t, params, self.input_transform, self.output_transform)
+            return self.eval_nn(t, params)
         if self.eq_type == "statio_PDE":
             (x, params) = args
-            return self._eval_nn(x, params, self.input_transform, self.output_transform)
+            return self.eval_nn(x, params)
         if self.eq_type == "nonstatio_PDE":
             (t, x, params) = args
             t_x = jnp.concatenate([t, x], axis=-1)
-            return self._eval_nn(
-                t_x, params, self.input_transform, self.output_transform
-            )
+            return self.eval_nn(t_x, params)
         raise ValueError("Wrong value for self.eq_type")
 
-    def _eval_nn(self, inputs, params, input_transform, output_transform):
+    def eval_nn(
+        self,
+        inputs: Float[Array, "input_dim"],
+        params: Params | PyTree,
+    ) -> Float[Array, "output_dim"]:
         """
-        inner function to factorize code. apply_fn (which takes varying forms)
-        call _eval_nn which always have the same content.
+        Evaluate the PINN on some inputs with some params.
         """
         try:
-            model = eqx.combine(params["nn_params"], self.static)
-        except (KeyError, TypeError) as e:  # give more flexibility
+            model = eqx.combine(params.nn_params, self.static)
+        except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
             model = eqx.combine(params, self.static)
-        res = output_transform(inputs, model(input_transform(inputs, params)).squeeze())
+        res = self.output_transform(
+            inputs, model(self.input_transform(inputs, params)).squeeze(), params
+        )
 
         if self.output_slice is not None:
             res = res[self.output_slice]
@@ -128,15 +180,20 @@ class PINN(eqx.Module):
 
 
 def create_PINN(
-    key,
-    eqx_list,
-    eq_type,
-    dim_x=0,
-    input_transform=None,
-    output_transform=None,
-    shared_pinn_outputs=None,
-    slice_solution=None,
-):
+    key: Key,
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+    dim_x: int = 0,
+    input_transform: Callable[
+        [Float[Array, "input_dim"], Params], Float[Array, "output_dim"]
+    ] = None,
+    output_transform: Callable[
+        [Float[Array, "input_dim"], Float[Array, "output_dim"], Params],
+        Float[Array, "output_dim"],
+    ] = None,
+    shared_pinn_outputs: tuple[slice] = None,
+    slice_solution: slice = None,
+) -> PINN | list[PINN]:
     r"""
     Utility function to create a standard PINN neural network with the equinox
     library.
@@ -144,22 +201,25 @@ def create_PINN(
     Parameters
     ----------
     key
-        A jax random key that will be used to initialize the network parameters
+        A JAX random key that will be used to initialize the network
+        parameters.
     eqx_list
-        A list of list of successive equinox modules and activation functions to
-        describe the PINN architecture. The inner lists have the eqx module or
-        axtivation function as first item, other items represents arguments
-        that could be required (eg. the size of the layer).
-        __Note:__ the `key` argument need not be given.
-        Thus typical example is `eqx_list=
-        [[eqx.nn.Linear, 2, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 20],
-        [jax.nn.tanh],
-        [eqx.nn.Linear, 20, 1]
-        ]`
+        A tuple of tuples of successive equinox modules and activation
+        functions to describe the PINN architecture. The inner tuples must have
+        the eqx module or activation function as first item, other items
+        represent arguments that could be required (eg. the size of the layer).
+
+        The `key` argument do not need to be given.
+
+        A typical example is `eqx_list = (
+            (eqx.nn.Linear, input_dim, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh,),
+            (eqx.nn.Linear, 20, output_dim)
+        )`.
     eq_type
         A string with three possibilities.
         "ODE": the PINN is called with one input `t`.
@@ -169,17 +229,19 @@ def create_PINN(
         can be high dimensional.
         **Note**: the input dimension as given in eqx_list has to match the sum
         of the dimension of `t` + the dimension of `x` or the output dimension
-        after the `input_transform` function
+        after the `input_transform` function.
     dim_x
-        An integer. The dimension of `x`. Default `0`
+        An integer. The dimension of `x`. Default `0`.
     input_transform
         A function that will be called before entering the PINN. Its output(s)
-        must match the PINN inputs. Its inputs are the PINN inputs (`t` and/or
-        `x` concatenated together and the parameters). Default is the No operation
+        must match the PINN inputs (except for the parameters).
+        Its inputs are the PINN inputs (`t` and/or `x` concatenated together)
+        and the parameters. Default is no operation.
     output_transform
-        A function with arguments the same input(s) as the PINN AND the PINN
-        output that will be called after exiting the PINN. Default is the No
-        operation
+        A function with arguments begin the same input as the PINN, the PINN
+        output and the parameter. This function will be called after exiting
+        the PINN.
+        Default is no operation.
     shared_pinn_outputs
         Default is None, for a stantard PINN.
         A tuple of jnp.s\_[] (slices) to determine the different output for each
@@ -192,15 +254,18 @@ def create_PINN(
     slice_solution
         A jnp.s\_ object which indicates which axis of the PINN output is
         dedicated to the actual equation solution. Default None
-        means that slice_solution = the whole PINN output. This argument is useful
-        when the PINN is also used to output equation parameters for example
-        Note that it must be a slice and not an integer (a preprocessing of the
-        user provided argument takes care of it)
+        means that slice_solution = the whole PINN output. This argument is
+        useful when the PINN is also used to output equation parameters for
+        example Note that it must be a slice and not an integer (a
+        preprocessing of the user provided argument takes care of it).
 
 
     Returns
     -------
-    `u`, a :class:`.PINN` object which inherits from `eqx.Module` (hence callable). This comes with a bound method :func:`u.init_params() <PINN.init_params>`. When `shared_pinn_ouput` is not None, a list of :class:`.PINN` with the same structure is returned, only differing by there final slicing of the network output.
+    pinn
+        A PINN instance or, when `shared_pinn_ouput` is not None,
+        a list of PINN instances with the same structure is returned,
+        only differing by there final slicing of the network output.
 
     Raises
     ------
@@ -240,23 +305,30 @@ def create_PINN(
 
     if output_transform is None:
 
-        def output_transform(_in_pinn, _out_pinn):
+        def output_transform(_in_pinn, _out_pinn, _params):
             return _out_pinn
 
-    mlp = _MLP(key, eqx_list)
+    mlp = _MLP(key=key, eqx_list=eqx_list)
 
     if shared_pinn_outputs is not None:
         pinns = []
         for output_slice in shared_pinn_outputs:
             pinn = PINN(
-                mlp,
-                slice_solution,
-                eq_type,
-                input_transform,
-                output_transform,
-                output_slice,
+                mlp=mlp,
+                slice_solution=slice_solution,
+                eq_type=eq_type,
+                input_transform=input_transform,
+                output_transform=output_transform,
+                output_slice=output_slice,
             )
             pinns.append(pinn)
         return pinns
-    pinn = PINN(mlp, slice_solution, eq_type, input_transform, output_transform, None)
+    pinn = PINN(
+        mlp=mlp,
+        slice_solution=slice_solution,
+        eq_type=eq_type,
+        input_transform=input_transform,
+        output_transform=output_transform,
+        output_slice=None,
+    )
     return pinn

jinns/utils/_save_load.py

@@ -2,58 +2,64 @@
 Implements save and load functions
 """
 
+from typing import Callable, Literal
 import pickle
 import jax
 import equinox as eqx
 
-from jinns.utils._pinn import create_PINN
-from jinns.utils._spinn import create_SPINN
-from jinns.utils._hyperpinn import create_HYPERPINN
+from jinns.utils._pinn import create_PINN, PINN
+from jinns.utils._spinn import create_SPINN, SPINN
+from jinns.utils._hyperpinn import create_HYPERPINN, HYPERPINN
+from jinns.parameters._params import Params, ParamsDict
 
 
-def function_to_string(eqx_list):
+def function_to_string(
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...]
+) -> tuple[tuple[str, int, int] | str, ...]:
     """
     We need this transformation for eqx_list to be pickled
 
-    From `[[eqx.nn.Linear, 2, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 1]` to
-    `[["Linear", 2, 20],
-                ["tanh"],
-                ["Linear", 20, 20],
-                ["tanh"],
-                ["Linear", 20, 20],
-                ["tanh"],
-                ["Linear", 20, 1]`
+    From `((eqx.nn.Linear, 2, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 1))` to
+    `(("Linear", 2, 20),
+                ("tanh"),
+                ("Linear", 20, 20),
+                ("tanh"),
+                ("Linear", 20, 20),
+                ("tanh"),
+                ("Linear", 20, 1))`
     """
     return jax.tree_util.tree_map(
         lambda x: x.__name__ if hasattr(x, "__call__") else x, eqx_list
     )
 
 
-def string_to_function(eqx_list_with_string):
+def string_to_function(
+    eqx_list_with_string: tuple[tuple[str, int, int] | str, ...]
+) -> tuple[tuple[Callable, int, int] | Callable, ...]:
     """
     We need this transformation for eqx_list at the loading ("unpickling")
     operation.
 
-    From `[["Linear", 2, 20],
-                ["tanh"],
-                ["Linear", 20, 20],
-                ["tanh"],
-                ["Linear", 20, 20],
-                ["tanh"],
-                ["Linear", 20, 1]`
-    to  `[[eqx.nn.Linear, 2, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 20],
-            [jax.nn.tanh],
-            [eqx.nn.Linear, 20, 1]` to
+    From `(("Linear", 2, 20),
+                ("tanh"),
+                ("Linear", 20, 20),
+                ("tanh"),
+                ("Linear", 20, 20),
+                ("tanh"),
+                ("Linear", 20, 1))`
+    to  `((eqx.nn.Linear, 2, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 20),
+            (jax.nn.tanh),
+            (eqx.nn.Linear, 20, 1))`
     """
 
     def _str_to_fun(l):
@@ -76,16 +82,36 @@ def string_to_function(eqx_list_with_string):
     )
 
 
-def save_pinn(filename, u, params, kwargs_creation):
+def save_pinn(
+    filename: str,
+    u: PINN | HYPERPINN | SPINN,
+    params: Params | ParamsDict,
+    kwargs_creation,
+):
     """
     Save a PINN / HyperPINN / SPINN model
     This function creates 3 files, beggining by `filename`
 
      1. an eqx file to save the eqx.Module (the PINN, HyperPINN, ...)
-     2. a pickle file for the parameters
-     3. a pickle file for the arguments that have been used at PINN
-
-     creation and that we need to reconstruct the eqx.module later on.
+     2. a pickle file for the parameters of the equation
+     3. a pickle file for the arguments that have been used at PINN creation
+     and that we need to reconstruct the eqx.module later on.
+
+    Note that the equation parameters `Params.eq_params` go in the
+    pickle file while the neural network parameters `Params.nn_params` go in
+    the `"*-module.eqx"` file (normal behaviour with `eqx.
+    tree_serialise_leaves`).
+
+    Equation parameters are saved apart because the initial type of attribute
+    `params` in PINN / HYPERPINN / SPINN is not `Params` nor `ParamsDict`
+    but `PyTree` as inherited from `eqx.partition`.
+    Therefore, if we want to ensure a proper serialization/deserialization:
+    - we cannot save a `Params` object at this
+      attribute field ; the `Params` object must be split into `Params.nn_params`
+      (type `PyTree`) and `Params.eq_params` (type `dict`).
+    - in the case of a `ParamsDict` we cannot save `ParamsDict.nn_params` at
+      the attribute field `params` because it is not a `PyTree` (as expected in
+      the PINN / HYPERPINN / SPINN signature) but it is still a dictionary.
 
     Parameters
     ----------
@@ -94,17 +120,32 @@ def save_pinn(filename, u, params, kwargs_creation):
     u
         The PINN
     params
-        The dictionary of parameters of the model.
-        Typically, it is a dictionary of
-        dictionaries: `eq_params` and `nn_params`, respectively the
-        differential equation parameters and the neural network parameter
+        Params or ParamsDict to be save
     kwargs_creation
         The dictionary of arguments that were used to create the PINN, e.g.
         the layers list, O/PDE type, etc.
     """
-    eqx.tree_serialise_leaves(filename + "-module.eqx", u)
-    with open(filename + "-parameters.pkl", "wb") as f:
-        pickle.dump(params, f)
+    if isinstance(params, Params):
+        if isinstance(u, HYPERPINN):
+            u = eqx.tree_at(lambda m: m.params_hyper, u, params)
+        elif isinstance(u, (PINN, SPINN)):
+            u = eqx.tree_at(lambda m: m.params, u, params)
+        eqx.tree_serialise_leaves(filename + "-module.eqx", u)
+
+    elif isinstance(params, ParamsDict):
+        for key, params_ in params.nn_params.items():
+            if isinstance(u, HYPERPINN):
+                u = eqx.tree_at(lambda m: m.params_hyper, u, params_)
+            elif isinstance(u, (PINN, SPINN)):
+                u = eqx.tree_at(lambda m: m.params, u, params_)
+            eqx.tree_serialise_leaves(filename + f"-module_{key}.eqx", u)
+
+    else:
+        raise ValueError("The parameters to be saved must be a Params or a ParamsDict")
+
+    with open(filename + "-eq_params.pkl", "wb") as f:
+        pickle.dump(params.eq_params, f)
+
     kwargs_creation = kwargs_creation.copy()  # avoid side-effect that would be
     # very probably harmless anyway
 
@@ -124,7 +165,11 @@ def save_pinn(filename, u, params, kwargs_creation):
         pickle.dump(kwargs_creation, f)
 
 
-def load_pinn(filename, type_):
+def load_pinn(
+    filename: str,
+    type_: Literal["pinn", "hyperpinn", "spinn"],
+    key_list_for_paramsdict: list[str] = None,
+) -> tuple[eqx.Module, Params | ParamsDict]:
     """
     Load a PINN model. This function needs to access 3 files :
     `{filename}-module.eqx`, `{filename}-parameters.pkl` and
@@ -132,27 +177,35 @@ def load_pinn(filename, type_):
 
     These files are created by `jinns.utils.save_pinn`.
 
-    Note that this requires equinox v0.11.3 (currently latest version) for the
+    Note that this requires equinox>v0.11.3 for the
     `eqx.filter_eval_shape` to work.
 
+    See note in `save_pinn` for more details about the saving process
+
     Parameters
     ----------
     filename
         Filename (prefix) without extension.
     type_
         Type of model to load. Must be in ["pinn", "hyperpinn", "spinn"].
+    key_list_for_paramsdict
+        Pass the name of the keys of the dictionnary `ParamsDict.nn_params`. Default is None. In this case, we expect to retrieve a ParamsDict.
 
     Returns
     -------
     u_reloaded
         The reloaded PINN
-    params_reloaded
+    params
         The reloaded parameters
     """
     with open(filename + "-arguments.pkl", "rb") as f:
         kwargs_reloaded = pickle.load(f)
-    with open(filename + "-parameters.pkl", "rb") as f:
-        params_reloaded = pickle.load(f)
+    try:
+        with open(filename + "-eq_params.pkl", "rb") as f:
+            eq_params_reloaded = pickle.load(f)
+    except FileNotFoundError:
+        eq_params_reloaded = {}
+        print("No pickle file for equation parameters found!")
     kwargs_reloaded["eqx_list"] = string_to_function(kwargs_reloaded["eqx_list"])
     if type_ == "pinn":
         # next line creates a shallow model, the jax arrays are just shapes and
@@ -167,9 +220,21 @@ def load_pinn(filename, type_):
         u_reloaded_shallow = eqx.filter_eval_shape(create_HYPERPINN, **kwargs_reloaded)
     else:
         raise ValueError(f"{type_} is not valid")
-    # now the empty structure is populated with the actual saved array values
-    # stored in the eqx file
-    u_reloaded = eqx.tree_deserialise_leaves(
-        filename + "-module.eqx", u_reloaded_shallow
-    )
-    return u_reloaded, params_reloaded
+    if key_list_for_paramsdict is None:
+        # now the empty structure is populated with the actual saved array values
+        # stored in the eqx file
+        u_reloaded = eqx.tree_deserialise_leaves(
+            filename + "-module.eqx", u_reloaded_shallow
+        )
+        params = Params(
+            nn_params=u_reloaded.init_params(), eq_params=eq_params_reloaded
+        )
+    else:
+        nn_params_dict = {}
+        for key in key_list_for_paramsdict:
+            u_reloaded = eqx.tree_deserialise_leaves(
+                filename + f"-module_{key}.eqx", u_reloaded_shallow
+            )
+            nn_params_dict[key] = u_reloaded.init_params()
+        params = ParamsDict(nn_params=nn_params_dict, eq_params=eq_params_reloaded)
+    return u_reloaded, params