jinns 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

jinns/nn/_mlp.py

@@ -2,16 +2,30 @@
 Implements utility function to create PINNs
 """
 
-from typing import Callable, Literal, Self, Union, Any
+from __future__ import annotations
+
+from typing import Callable, Literal, Self, Union, Any, TYPE_CHECKING, cast
 from dataclasses import InitVar
 import jax
 import equinox as eqx
-
+from typing import Protocol
 from jaxtyping import Array, Key, PyTree, Float
 
 from jinns.parameters._params import Params
 from jinns.nn._pinn import PINN
 
+if TYPE_CHECKING:
+
+    class CallableMLPModule(Protocol):
+        """
+        Basically just a way to add a __call__ to an eqx.Module.
+        https://github.com/patrick-kidger/equinox/issues/1002
+        We chose the strutural subtyping of protocols instead of subclassing an
+        eqx.Module just to add a __call__ here
+        """
+
+        def __call__(self, *_, **__) -> Array: ...
+
 
 class MLP(eqx.Module):
     """
@@ -21,7 +35,7 @@ class MLP(eqx.Module):
     ----------
     key : InitVar[Key]
         A jax random key for the layer initializations.
-    eqx_list : InitVar[tuple[tuple[Callable, int, int] | Callable, ...]]
+    eqx_list : InitVar[tuple[tuple[Callable, int, int] | tuple[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
@@ -29,23 +43,23 @@ class MLP(eqx.Module):
         The `key` argument need not be given.
         Thus typical example is `eqx_list=
         ((eqx.nn.Linear, 2, 20),
-            jax.nn.tanh,
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
+            (jax.nn.tanh,),
             (eqx.nn.Linear, 20, 1)
         )`.
     """
 
     key: InitVar[Key] = eqx.field(kw_only=True)
-    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
-        kw_only=True
+    eqx_list: InitVar[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] = (
+        eqx.field(kw_only=True)
     )
 
     # 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)
+    layers: list[CallableMLPModule | Callable[[Array], Array]] = eqx.field(init=False)
 
     def __post_init__(self, key, eqx_list):
         self.layers = []
@@ -63,7 +77,7 @@ class MLP(eqx.Module):
                 self.layers.append(l[0](*l[1:], key=subkey))
                 k += 1
 
-    def __call__(self, t: Float[Array, "input_dim"]) -> Float[Array, "output_dim"]:
+    def __call__(self, t: Float[Array, "  input_dim"]) -> Float[Array, "  output_dim"]:
         for layer in self.layers:
             t = layer(t)
         return t
@@ -81,19 +95,30 @@ class PINN_MLP(PINN):
     def create(
         cls,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
-        eqx_network: eqx.nn.MLP = None,
+        eqx_network: eqx.nn.MLP | MLP | None = None,
         key: Key = None,
-        eqx_list: tuple[tuple[Callable, int, int] | Callable, ...] = None,
-        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,
-        slice_solution: slice = None,
+        eqx_list: tuple[tuple[Callable, int, int] | tuple[Callable], ...] | None = None,
+        input_transform: (
+            Callable[
+                [Float[Array, "  input_dim"], Params[Array]],
+                Float[Array, "  output_dim"],
+            ]
+            | None
+        ) = None,
+        output_transform: (
+            Callable[
+                [
+                    Float[Array, "  input_dim"],
+                    Float[Array, "  output_dim"],
+                    Params[Array],
+                ],
+                Float[Array, "  output_dim"],
+            ]
+            | None
+        ) = None,
+        slice_solution: slice | None = None,
         filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = None,
-    ) -> tuple[Self, PyTree]:
+    ) -> tuple[Self, PINN]:
         r"""
         Instanciate standard PINN MLP object. The actual NN is either passed as
         a eqx.nn.MLP (`eqx_network` argument) or constructed as a custom
@@ -179,14 +204,14 @@ class PINN_MLP(PINN):
                 raise ValueError(
                     "If eqx_network is None, then key and eqx_list must be provided"
                 )
-            eqx_network = MLP(key=key, eqx_list=eqx_list)
+            eqx_network = cast(MLP, MLP(key=key, eqx_list=eqx_list))
 
         mlp = cls(
             eqx_network=eqx_network,
-            slice_solution=slice_solution,
+            slice_solution=slice_solution,  # type: ignore
             eq_type=eq_type,
-            input_transform=input_transform,
-            output_transform=output_transform,
+            input_transform=input_transform,  # type: ignore
+            output_transform=output_transform,  # type: ignore
             filter_spec=filter_spec,
         )
         return mlp, mlp.init_params

jinns/nn/_pinn.py

@@ -2,15 +2,19 @@
 Implement abstract class for PINN architectures
 """
 
-from typing import Literal, Callable, Union, Any
+from __future__ import annotations
+
+from typing import Callable, Union, Any, Literal, overload
 from dataclasses import InitVar
 import equinox as eqx
 from jaxtyping import Float, Array, PyTree
 import jax.numpy as jnp
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+from jinns.nn._utils import _PyTree_to_Params
 
 
-class PINN(eqx.Module):
+class PINN(AbstractPINN):
     r"""
     Base class for PINN objects. It can be seen as a wrapper on
     an `eqx.Module` which actually implement the NN architectures, with extra
@@ -57,12 +61,12 @@ class PINN(eqx.Module):
         **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"]]
+    input_transform : Callable[[Float[Array, " input_dim"], Params[Array]], 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"]]
+    output_transform : Callable[[Float[Array, " input_dim"], Float[Array, " output_dim"], Params[Array]], 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.
@@ -84,16 +88,16 @@ class PINN(eqx.Module):
         "nonstatio_PDE"]`
     """
 
-    slice_solution: slice = eqx.field(static=True, kw_only=True, default=None)
     eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
         static=True, kw_only=True
     )
+    slice_solution: slice = eqx.field(static=True, kw_only=True, default=None)
     input_transform: Callable[
-        [Float[Array, "input_dim"], Params], Float[Array, "output_dim"]
+        [Float[Array, " input_dim"], Params[Array]], Float[Array, " output_dim"]
     ] = eqx.field(static=True, kw_only=True, default=None)
     output_transform: Callable[
-        [Float[Array, "input_dim"], Float[Array, "output_dim"], Params],
-        Float[Array, "output_dim"],
+        [Float[Array, " input_dim"], Float[Array, " output_dim"], Params[Array]],
+        Float[Array, " output_dim"],
     ] = eqx.field(static=True, kw_only=True, default=None)
 
     eqx_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
@@ -101,11 +105,10 @@ class PINN(eqx.Module):
         static=True, kw_only=True, default=eqx.is_inexact_array
     )
 
-    init_params: PyTree = eqx.field(init=False)
-    static: PyTree = eqx.field(init=False, static=True)
+    init_params: PINN = eqx.field(init=False)
+    static: PINN = eqx.field(init=False, static=True)
 
     def __post_init__(self, eqx_network):
-
         if self.eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
             raise RuntimeError("Wrong parameter value for eq_type")
         # saving the static part of the model and initial parameters
@@ -154,18 +157,32 @@ class PINN(eqx.Module):
 
         return network(inputs)
 
+    @overload
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        inputs: Float[Array, " input_dim"],
+        params: PyTree,
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
     def __call__(
         self,
-        inputs: Float[Array, "input_dim"],
-        params: Params | ParamsDict | PyTree,
+        inputs: Float[Array, " input_dim"],
+        params: Params[Array],
         *args,
         **kwargs,
-    ) -> Float[Array, "output_dim"]:
+    ) -> Float[Array, " output_dim"]:
         """
         A proper __call__ implementation performs an eqx.combine here with
         `params` and `self.static` to recreate the callable eqx.Module
         architecture. The rest of the content of this function is dependent on
         the network.
+
+        Note that that thanks to the decorator, params can also directly be the
+        PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
         """
 
         if len(inputs.shape) == 0:
@@ -174,10 +191,7 @@ class PINN(eqx.Module):
             # DataGenerators)
             inputs = inputs[None]
 
-        try:
-            model = eqx.combine(params.nn_params, self.static)
-        except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
-            model = eqx.combine(params, self.static)
+        model = eqx.combine(params.nn_params, self.static)
 
         # evaluate the model
         res = self.eval(model, self.input_transform(inputs, params), *args, **kwargs)

jinns/nn/_ppinn.py

@@ -2,17 +2,20 @@
 Implements utility function to create PINNs
 """
 
-from typing import Callable, Literal, Self
+from __future__ import annotations
+
+from typing import Callable, Literal, Self, cast, overload
 from dataclasses import InitVar
 import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-from jaxtyping import Array, Key, PyTree, Float
+from jaxtyping import Array, Key, Float, PyTree
 
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
 from jinns.nn._pinn import PINN
 from jinns.nn._mlp import MLP
+from jinns.nn._utils import _PyTree_to_Params
 
 
 class PPINN_MLP(PINN):
@@ -39,12 +42,12 @@ class PPINN_MLP(PINN):
         **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"]]
+    input_transform : Callable[[Float[Array, " input_dim"], Params[Array]], Float[Array, " output_dim"]]
         A function that will be called before entering the PPINN. Its output(s)
         must match the PPINN inputs (except for the parameters).
         Its inputs are the PPINN 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"]]
+    output_transform : Callable[[Float[Array, " input_dim"], Float[Array, " output_dim"], Params[Array]], Float[Array, " output_dim"]]
         A function with arguments begin the same input as the PPINN, the PPINN
         output and the parameter. This function will be called after exiting
         the PPINN.
@@ -63,25 +66,46 @@ class PPINN_MLP(PINN):
     """
 
     eqx_network_list: InitVar[list[eqx.Module]] = eqx.field(kw_only=True)
+    init_params: tuple[PINN, ...] = eqx.field(
+        init=False
+    )  # overriding parent attribute type
+    static: tuple[PINN, ...] = eqx.field(
+        init=False, static=True
+    )  # overriding parent attribute type
 
     def __post_init__(self, eqx_network, eqx_network_list):
         super().__post_init__(
             eqx_network=eqx_network_list[0],  # this is not used since it is
             # overwritten just below
         )
-        self.init_params, self.static = (), ()
-        for eqx_network_ in eqx_network_list:
+        params, static = eqx.partition(eqx_network_list[0], self.filter_spec)
+        self.init_params, self.static = (params,), (static,)
+        for eqx_network_ in eqx_network_list[1:]:
             params, static = eqx.partition(eqx_network_, self.filter_spec)
             self.init_params = self.init_params + (params,)
             self.static = self.static + (static,)
 
+    @overload
+    @_PyTree_to_Params
     def __call__(
         self,
-        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
+        inputs: Float[Array, " input_dim"],
         params: PyTree,
-    ) -> Float[Array, "output_dim"]:
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        inputs: Float[Array, " 1"] | Float[Array, " dim"] | Float[Array, " 1+dim"],
+        params: Params[Array],
+    ) -> Float[Array, " output_dim"]:
         """
         Evaluate the PPINN on some inputs with some params.
+
+        Note that that thanks to the decorator, params can also directly be the
+        PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
         """
         if len(inputs.shape) == 0:
             # This can happen often when the user directly provides some
@@ -92,14 +116,14 @@ class PPINN_MLP(PINN):
 
         outs = []
 
-        try:
-            for params_, static in zip(params.nn_params, self.static):
-                model = eqx.combine(params_, static)
-                outs += [model(transformed_inputs)]
-        except (KeyError, AttributeError, TypeError) as e:
-            for params_, static in zip(params, self.static):
-                model = eqx.combine(params_, static)
-                outs += [model(transformed_inputs)]
+        # try:
+        for params_, static in zip(params.nn_params, self.static):
+            model = eqx.combine(params_, static)
+            outs += [model(transformed_inputs)]  # type: ignore
+        # except (KeyError, AttributeError, TypeError) as e:
+        #    for params_, static in zip(params, self.static):
+        #        model = eqx.combine(params_, static)
+        #        outs += [model(transformed_inputs)]
         # Note that below is then a global output transform
         res = self.output_transform(inputs, jnp.concatenate(outs, axis=0), params)
 
@@ -112,18 +136,31 @@ class PPINN_MLP(PINN):
     def create(
         cls,
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
-        eqx_network_list: list[eqx.nn.MLP] = None,
+        eqx_network_list: list[eqx.nn.MLP | MLP] | None = None,
         key: Key = None,
-        eqx_list_list: list[tuple[tuple[Callable, int, int] | Callable, ...]] = None,
-        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,
-        slice_solution: slice = None,
-    ) -> tuple[Self, PyTree]:
+        eqx_list_list: (
+            list[tuple[tuple[Callable, int, int] | tuple[Callable], ...]] | None
+        ) = None,
+        input_transform: (
+            Callable[
+                [Float[Array, " input_dim"], Params[Array]],
+                Float[Array, " output_dim"],
+            ]
+            | None
+        ) = None,
+        output_transform: (
+            Callable[
+                [
+                    Float[Array, " input_dim"],
+                    Float[Array, " output_dim"],
+                    Params[Array],
+                ],
+                Float[Array, " output_dim"],
+            ]
+            | None
+        ) = None,
+        slice_solution: slice | None = None,
+    ) -> tuple[Self, tuple[PINN, ...]]:
         r"""
         Utility function to create a Parrallel PINN neural network for Jinns.
 
@@ -189,15 +226,14 @@ class PPINN_MLP(PINN):
             eqx_network_list = []
             for eqx_list in eqx_list_list:
                 key, subkey = jax.random.split(key, 2)
-                print(subkey)
                 eqx_network_list.append(MLP(key=subkey, eqx_list=eqx_list))
 
         ppinn = cls(
-            eqx_network=None,
-            eqx_network_list=eqx_network_list,
-            slice_solution=slice_solution,
+            eqx_network=None,  # type: ignore
+            eqx_network_list=cast(list[eqx.Module], eqx_network_list),
+            slice_solution=slice_solution,  # type: ignore
             eq_type=eq_type,
-            input_transform=input_transform,
-            output_transform=output_transform,
+            input_transform=input_transform,  # type: ignore
+            output_transform=output_transform,  # type: ignore
         )
         return ppinn, ppinn.init_params

jinns/nn/_save_load.py

@@ -12,7 +12,7 @@ from jinns.nn._spinn import SPINN
 from jinns.nn._mlp import PINN_MLP
 from jinns.nn._spinn_mlp import SPINN_MLP
 from jinns.nn._hyperpinn import HyperPINN
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
 
 
 def function_to_string(
@@ -87,7 +87,7 @@ def string_to_function(
 def save_pinn(
     filename: str,
     u: PINN | HyperPINN | SPINN,
-    params: Params | ParamsDict,
+    params: Params,
     kwargs_creation,
 ):
     """
@@ -105,15 +105,12 @@ def save_pinn(
     tree_serialise_leaves`).
 
     Equation parameters are saved apart because the initial type of attribute
-    `params` in PINN / HyperPINN / SPINN is not `Params` nor `ParamsDict`
+    `params` in PINN / HyperPINN / SPINN is not `Params`
     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
     ----------
@@ -122,28 +119,16 @@ def save_pinn(
     u
         The PINN
     params
-        Params or ParamsDict to be save
+        Params to be saved
     kwargs_creation
         The dictionary of arguments that were used to create the PINN, e.g.
         the layers list, O/PDE type, etc.
     """
-    if isinstance(params, Params):
-        if isinstance(u, HyperPINN):
-            u = eqx.tree_at(lambda m: m.init_params_hyper, u, params)
-        elif isinstance(u, (PINN, SPINN)):
-            u = eqx.tree_at(lambda m: m.init_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.init_params_hyper, u, params_)
-            elif isinstance(u, (PINN, SPINN)):
-                u = eqx.tree_at(lambda m: m.init_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")
+    if isinstance(u, HyperPINN):
+        u = eqx.tree_at(lambda m: m.init_params_hyper, u, params)
+    elif isinstance(u, (PINN, SPINN)):
+        u = eqx.tree_at(lambda m: m.init_params, u, params)
+    eqx.tree_serialise_leaves(filename + "-module.eqx", u)
 
     with open(filename + "-eq_params.pkl", "wb") as f:
         pickle.dump(params.eq_params, f)
@@ -170,8 +155,7 @@ def save_pinn(
 def load_pinn(
     filename: str,
     type_: Literal["pinn_mlp", "hyperpinn", "spinn_mlp"],
-    key_list_for_paramsdict: list[str] = None,
-) -> tuple[eqx.Module, Params | ParamsDict]:
+) -> tuple[eqx.Module, Params]:
     """
     Load a PINN model. This function needs to access 3 files :
     `{filename}-module.eqx`, `{filename}-parameters.pkl` and
@@ -190,8 +174,6 @@ def load_pinn(
         Filename (prefix) without extension.
     type_
         Type of model to load. Must be in ["pinn_mlp", "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
     -------
@@ -228,29 +210,17 @@ def load_pinn(
         )
     else:
         raise ValueError(f"{type_} is not valid")
-    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
+    # 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
+    )
+    if isinstance(u_reloaded, HyperPINN):
+        params = Params(
+            nn_params=u_reloaded.init_params_hyper, eq_params=eq_params_reloaded
         )
-        if isinstance(u_reloaded, HyperPINN):
-            params = Params(
-                nn_params=u_reloaded.init_params_hyper, eq_params=eq_params_reloaded
-            )
-        elif isinstance(u_reloaded, (PINN, SPINN)):
-            params = Params(
-                nn_params=u_reloaded.init_params, eq_params=eq_params_reloaded
-            )
+    elif isinstance(u_reloaded, (PINN, SPINN)):
+        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
-            )
-            if isinstance(u_reloaded, HyperPINN):
-                nn_params_dict[key] = u_reloaded.init_params_hyper
-            elif isinstance(u_reloaded, (PINN, SPINN)):
-                nn_params_dict[key] = u_reloaded.init_params
-        params = ParamsDict(nn_params=nn_params_dict, eq_params=eq_params_reloaded)
+        raise ValueError("Wrong type for u_reloaded")
     return u_reloaded, params

jinns/nn/_spinn.py

@@ -1,14 +1,17 @@
-from typing import Union, Callable, Any
+from __future__ import annotations
+from typing import Union, Callable, Any, Literal, overload
 from dataclasses import InitVar
 from jaxtyping import PyTree, Float, Array
 import jax
 import jax.numpy as jnp
 import equinox as eqx
 
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+from jinns.nn._utils import _PyTree_to_Params
 
 
-class SPINN(eqx.Module):
+class SPINN(AbstractPINN):
     """
     A Separable PINN object compatible with the rest of jinns.
 
@@ -47,21 +50,21 @@ class SPINN(eqx.Module):
 
     """
 
+    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"] = eqx.field(
+        static=True, kw_only=True
+    )
     d: int = eqx.field(static=True, kw_only=True)
     r: int = eqx.field(static=True, kw_only=True)
-    eq_type: str = eqx.field(static=True, kw_only=True)
     m: int = eqx.field(static=True, kw_only=True, default=1)
-
     filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = eqx.field(
         static=True, kw_only=True, default=None
     )
     eqx_spinn_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
 
-    init_params: PyTree = eqx.field(init=False)
-    static: PyTree = eqx.field(init=False, static=True)
+    init_params: SPINN = eqx.field(init=False)
+    static: SPINN = eqx.field(init=False, static=True)
 
     def __post_init__(self, eqx_spinn_network):
-
         if self.filter_spec is None:
             self.filter_spec = eqx.is_inexact_array
 
@@ -69,20 +72,34 @@ class SPINN(eqx.Module):
             eqx_spinn_network, self.filter_spec
         )
 
+    @overload
+    @_PyTree_to_Params
     def __call__(
         self,
-        t_x: Float[Array, "batch_size 1+dim"],
-        params: Params | ParamsDict | PyTree,
-    ) -> Float[Array, "output_dim"]:
+        inputs: Float[Array, " input_dim"],
+        params: PyTree,
+        *args,
+        **kwargs,
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        t_x: Float[Array, "  batch_size 1+dim"],
+        params: Params[Array],
+    ) -> Float[Array, "  output_dim"]:
         """
         Evaluate the SPINN on some inputs with some params.
+
+        Note that that thanks to the decorator, params can also directly be the
+        PyTree (SPINN, PINN_MLP, ...) that we get out of eqx.combine
         """
-        try:
-            spinn = eqx.combine(params.nn_params, self.static)
-        except (KeyError, AttributeError, TypeError) as e:
-            spinn = eqx.combine(params, self.static)
+        # try:
+        spinn = eqx.combine(params.nn_params, self.static)
+        # except (KeyError, AttributeError, TypeError) as e:
+        #    spinn = eqx.combine(params, self.static)
         v_model = jax.vmap(spinn)
-        res = v_model(t_x)
+        res = v_model(t_x)  # type: ignore
 
         a = ", ".join([f"{chr(97 + d)}z" for d in range(res.shape[1])])
         b = "".join([f"{chr(97 + d)}" for d in range(res.shape[1])])