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

jinns/nn/_pinn.py

@@ -0,0 +1,190 @@
+"""
+Implement abstract class for PINN architectures
+"""
+
+from typing import Literal, Callable, Union, Any
+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
+
+
+class PINN(eqx.Module):
+    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
+    arguments handling the "physics-informed" aspect.
+
+    !!! Note
+        We use the `eqx.partition` and `eqx.combine` strategy of Equinox: a
+        `filter_spec` is applied on the PyTree and splits it into two PyTree with
+        the same structure: a static one (invisible to JAX transform such as JIT,
+        grad, etc.) and dynamic one. By convention, anything not static is
+        considered a parameter in Jinns.
+
+    For compatibility with jinns, we require that a `PINN` architecture:
+
+        1) has an eqx.Module (`eqx_network`) InitVar passed to __post_init__
+        representing the network architecture.
+        2) calls `eqx.partition` in __post_init__ in order to store the
+        static part of the model and the initial parameters.
+        3) has a `eq_type` argument, used for handling internal operations in
+        jinns.
+        4) has a `slice_solution` argument. It is 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.
+    For example, this argument is useful when the PINN is also used to output
+    equation parameters. Note that it must be a slice and not an integer (a
+    preprocessing of the user provided argument takes care of it).
+
+    Parameters
+    ----------
+    slice_solution : slice
+        Default is jnp.s\_[...]. 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.
+    eqx_network : eqx.Module
+        The actual neural network instanciated as an eqx.Module.
+    filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+        Default is `eqx.is_inexact_array`. This tells Jinns what to consider as
+        a trainable parameter. Quoting from equinox documentation:
+        a PyTree whose structure should be a prefix of the structure of pytree.
+        Each of its leaves should either be 1) True, in which case the leaf or
+        subtree is kept; 2) False, in which case the leaf or subtree is
+        replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+
+
+    Raises
+    ------
+    RuntimeError
+        If the parameter value for eq_type is not in `["ODE", "statio_PDE",
+        "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
+    )
+    input_transform: Callable[
+        [Float[Array, "input_dim"], Params], 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"],
+    ] = eqx.field(static=True, kw_only=True, default=None)
+
+    eqx_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
+    filter_spec: PyTree[Union[bool, Callable[[Any], bool]]] = eqx.field(
+        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)
+
+    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
+
+        if self.filter_spec is None:
+            self.filter_spec = eqx.is_inexact_array
+
+        self.init_params, self.static = eqx.partition(eqx_network, self.filter_spec)
+
+        if self.input_transform is None:
+            self.input_transform = lambda _in, _params: _in
+
+        if self.output_transform is None:
+            self.output_transform = lambda _in_pinn, _out_pinn, _params: _out_pinn
+
+        if self.slice_solution is None:
+            self.slice_solution = jnp.s_[:]
+
+        if isinstance(self.slice_solution, int):
+            # rewrite it as a slice to ensure that axis does not disappear when
+            # indexing
+            self.slice_solution = jnp.s_[self.slice_solution : self.slice_solution + 1]
+
+    def eval(self, network, inputs, *args, **kwargs):
+        """How to call your Equinox module `network`. The purpose of this method
+        is to give more flexibility : user should re-implement `eval`
+        when inheriting from `PINN` if they desire more flexibility on how to
+        evaluate the network.
+
+        Defaults to using `network.__call__(inputs)` but it could be more refined *e.g.* `network.anymethod(inputs)`.
+
+        Parameters
+        ----------
+        network : eqx.Module
+            Your neural network with the parameters set, usually returned by
+            `eqx.combine(self.static, current_params)`.
+        inputs : Array
+            The inputs, evetually transformed by `self.input_transformed` if
+            specified by the user.
+
+        Returns
+        -------
+        Array
+            The output
+        """
+
+        return network(inputs)
+
+    def __call__(
+        self,
+        inputs: Float[Array, "input_dim"],
+        params: Params | ParamsDict | PyTree,
+        *args,
+        **kwargs,
+    ) -> 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.
+        """
+
+        if len(inputs.shape) == 0:
+            # This can happen often when the user directly provides some
+            # collocation points (eg for plotting, whithout using
+            # 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)
+
+        # evaluate the model
+        res = self.eval(model, self.input_transform(inputs, params), *args, **kwargs)
+
+        res = self.output_transform(inputs, res.squeeze(), params)
+
+        # force (1,) output for non vectorial solution (consistency)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res

jinns/nn/_ppinn.py

@@ -0,0 +1,203 @@
+"""
+Implements utility function to create PINNs
+"""
+
+from typing import Callable, Literal, Self
+from dataclasses import InitVar
+import jax
+import jax.numpy as jnp
+import equinox as eqx
+
+from jaxtyping import Array, Key, PyTree, Float
+
+from jinns.parameters._params import Params, ParamsDict
+from jinns.nn._pinn import PINN
+from jinns.nn._mlp import MLP
+
+
+class PPINN_MLP(PINN):
+    r"""
+    A PPINN MLP (Parallel PINN with MLPs) object which mimicks the PFNN architecture from
+    DeepXDE. This is in fact a PINN MLP that encompasses several PINN MLPs internally.
+
+    Parameters
+    ----------
+    slice_solution : slice
+        A jnp.s\_ object which indicates which axis of the PPINN output is
+        dedicated to the actual equation solution. Default None
+        means that slice_solution = the whole PPINN 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 PPINN is called with one input `t`.
+        "statio_PDE": the PPINN is called with one input `x`, `x`
+        can be high dimensional.
+        "nonstatio_PDE": the PPINN 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 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"]]
+        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.
+        Default is no operation.
+    filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+        Default is `eqx.is_inexact_array`. This tells Jinns what to consider as
+        a trainable parameter. Quoting from equinox documentation:
+        a PyTree whose structure should be a prefix of the structure of pytree.
+        Each of its leaves should either be 1) True, in which case the leaf or
+        subtree is kept; 2) False, in which case the leaf or subtree is
+        replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+    eqx_network_list
+            A list of eqx.nn.MLP objects with same input
+            dimensions. They represent the parallel subnetworks of the PPIN MLP.
+            Their respective outputs are concatenated.
+    """
+
+    eqx_network_list: InitVar[list[eqx.Module]] = eqx.field(kw_only=True)
+
+    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_, self.filter_spec)
+            self.init_params = self.init_params + (params,)
+            self.static = self.static + (static,)
+
+    def __call__(
+        self,
+        inputs: Float[Array, "1"] | Float[Array, "dim"] | Float[Array, "1+dim"],
+        params: PyTree,
+    ) -> Float[Array, "output_dim"]:
+        """
+        Evaluate the PPINN on some inputs with some params.
+        """
+        if len(inputs.shape) == 0:
+            # This can happen often when the user directly provides some
+            # collocation points (eg for plotting, whithout using
+            # DataGenerators)
+            inputs = inputs[None]
+        transformed_inputs = self.input_transform(inputs, params)
+
+        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)]
+        # Note that below is then a global output transform
+        res = self.output_transform(inputs, jnp.concatenate(outs, axis=0), params)
+
+        ## force (1,) output for non vectorial solution (consistency)
+        if not res.shape:
+            return jnp.expand_dims(res, axis=-1)
+        return res
+
+    @classmethod
+    def create(
+        cls,
+        eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
+        eqx_network_list: list[eqx.nn.MLP] = 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]:
+        r"""
+        Utility function to create a Parrallel PINN neural network for Jinns.
+
+        Parameters
+        ----------
+        eq_type
+            A string with three possibilities.
+            "ODE": the PPINN MLP is called with one input `t`.
+            "statio_PDE": the PPINN MLP is called with one input `x`, `x`
+            can be high dimensional.
+            "nonstatio_PDE": the PPINN MLP 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.
+        eqx_network_list
+                Default is None. A list of eqx.nn.MLP objects with same input
+                dimensions. They represent the parallel subnetworks of the PPIN MLP.
+                Their respective outputs are concatenated.
+        key
+            Default is None. Must be provided with `eqx_list_list` if
+            `eqx_network_list` is not provided. A JAX random key that will be used
+            to initialize the networks parameters.
+        eqx_list_list
+            Default is None. Must be provided if `eqx_network_list` is not
+            provided. A list of `eqx_list` (see `PINN_MLP.create()`). The input dimension must be the
+            same for each sub-`eqx_list`. Then the parallel subnetworks can be
+            different. Their respective outputs are concatenated.
+        input_transform
+            A function that will be called before entering the PPINN MLP. Its output(s)
+            must match the PPINN MLP inputs (except for the parameters).
+            Its inputs are the PPINN MLP inputs (`t` and/or `x` concatenated together)
+            and the parameters. Default is no operation.
+        output_transform
+            This function will be called after exiting
+            the PPINN MLP, i.e., on the concatenated outputs of all parallel networks
+            Default is no operation.
+        slice_solution
+            A jnp.s\_ object which indicates which axis of the PPINN MLP output is
+            dedicated to the actual equation solution. Default None
+            means that slice_solution = the whole PPINN MLP output. This argument is
+            useful when the PPINN MLP 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
+        -------
+        ppinn
+            A PPINN MLP instance
+        ppinn.init_params
+            An initial set of parameters for the PPINN MLP
+
+        """
+
+        if eqx_network_list is None:
+            if eqx_list_list is None or key is None:
+                raise ValueError(
+                    "If eqx_network_list is None, then key and eqx_list_list"
+                    " must be provided"
+                )
+
+            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,
+            eq_type=eq_type,
+            input_transform=input_transform,
+            output_transform=output_transform,
+        )
+        return ppinn, ppinn.init_params

jinns/{utils → nn}/_save_load.py

@@ -7,14 +7,16 @@ import pickle
 import jax
 import equinox as eqx
 
-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.nn._pinn import PINN
+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
 
 
 def function_to_string(
-    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...]
+    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
 ) -> tuple[tuple[str, int, int] | str, ...]:
     """
     We need this transformation for eqx_list to be pickled
@@ -40,7 +42,7 @@ def function_to_string(
 
 
 def string_to_function(
-    eqx_list_with_string: tuple[tuple[str, int, int] | str, ...]
+    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")
@@ -84,7 +86,7 @@ def string_to_function(
 
 def save_pinn(
     filename: str,
-    u: PINN | HYPERPINN | SPINN,
+    u: PINN | HyperPINN | SPINN,
     params: Params | ParamsDict,
     kwargs_creation,
 ):
@@ -103,7 +105,7 @@ 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` 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
@@ -111,7 +113,7 @@ def save_pinn(
       (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.
+      the PINN / HyperPINN / SPINN signature) but it is still a dictionary.
 
     Parameters
     ----------
@@ -126,18 +128,18 @@ def save_pinn(
         the layers list, O/PDE type, etc.
     """
     if isinstance(params, Params):
-        if isinstance(u, HYPERPINN):
-            u = eqx.tree_at(lambda m: m.params_hyper, u, 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.params, u, params)
+            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.params_hyper, u, 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.params, u, params_)
+                u = eqx.tree_at(lambda m: m.init_params, u, params_)
             eqx.tree_serialise_leaves(filename + f"-module_{key}.eqx", u)
 
     else:
@@ -167,7 +169,7 @@ def save_pinn(
 
 def load_pinn(
     filename: str,
-    type_: Literal["pinn", "hyperpinn", "spinn"],
+    type_: Literal["pinn_mlp", "hyperpinn", "spinn_mlp"],
     key_list_for_paramsdict: list[str] = None,
 ) -> tuple[eqx.Module, Params | ParamsDict]:
     """
@@ -187,7 +189,7 @@ def load_pinn(
     filename
         Filename (prefix) without extension.
     type_
-        Type of model to load. Must be in ["pinn", "hyperpinn", "spinn"].
+        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.
 
@@ -207,18 +209,22 @@ def load_pinn(
         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":
+    if type_ == "pinn_mlp":
         # next line creates a shallow model, the jax arrays are just shapes and
         # not populated, this just recreates the correct pytree structure
-        u_reloaded_shallow, _ = eqx.filter_eval_shape(create_PINN, **kwargs_reloaded)
-    elif type_ == "spinn":
-        u_reloaded_shallow, _ = eqx.filter_eval_shape(create_SPINN, **kwargs_reloaded)
+        u_reloaded_shallow, _ = eqx.filter_eval_shape(
+            PINN_MLP.create, **kwargs_reloaded
+        )
+    elif type_ == "spinn_mlp":
+        u_reloaded_shallow, _ = eqx.filter_eval_shape(
+            SPINN_MLP.create, **kwargs_reloaded
+        )
     elif type_ == "hyperpinn":
         kwargs_reloaded["eqx_list_hyper"] = string_to_function(
             kwargs_reloaded["eqx_list_hyper"]
         )
         u_reloaded_shallow, _ = eqx.filter_eval_shape(
-            create_HYPERPINN, **kwargs_reloaded
+            HyperPINN.create, **kwargs_reloaded
         )
     else:
         raise ValueError(f"{type_} is not valid")
@@ -228,13 +234,23 @@ def load_pinn(
         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)
+        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
+            )
     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
+            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)
     return u_reloaded, params

jinns/nn/_spinn.py

@@ -0,0 +1,106 @@
+from typing import Union, Callable, Any
+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
+
+
+class SPINN(eqx.Module):
+    """
+    A Separable PINN object compatible with the rest of jinns.
+
+    Parameters
+    ----------
+    d : int
+        The number of dimensions to treat separately, including time `t` if
+        used for non-stationnary equations.
+    r : int
+        An integer. The dimension of the embedding.
+    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`.
+    m : int
+        The output dimension of the neural network. According to
+        the SPINN article, a total embedding dimension of `r*m` is defined. We
+        then sum groups of `r` embedding dimensions to compute each output.
+        Default is 1.
+    filter_spec : PyTree[Union[bool, Callable[[Any], bool]]]
+        Default is `eqx.is_inexact_array`. This tells Jinns what to consider as
+        a trainable parameter. Quoting from equinox documentation:
+        a PyTree whose structure should be a prefix of the structure of pytree.
+        Each of its leaves should either be 1) True, in which case the leaf or
+        subtree is kept; 2) False, in which case the leaf or subtree is
+        replaced with replace; 3) a callable Leaf -> bool, in which case this is evaluated on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
+    eqx_spinn_network : eqx.Module
+        The actual neural network instanciated as an eqx.Module. It should be
+        an architecture taking `d` inputs and returning `d` times an embedding
+        of dimension `r`*`m`. See the Separable PINN paper for more details.
+
+    """
+
+    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)
+
+    def __post_init__(self, eqx_spinn_network):
+
+        if self.filter_spec is None:
+            self.filter_spec = eqx.is_inexact_array
+
+        self.init_params, self.static = eqx.partition(
+            eqx_spinn_network, self.filter_spec
+        )
+
+    def __call__(
+        self,
+        t_x: Float[Array, "batch_size 1+dim"],
+        params: Params | ParamsDict | PyTree,
+    ) -> Float[Array, "output_dim"]:
+        """
+        Evaluate the SPINN on some inputs with some params.
+        """
+        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)
+
+        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])])
+        res = jnp.stack(
+            [
+                jnp.einsum(
+                    f"{a} -> {b}",
+                    *(
+                        res[:, d, m * self.r : (m + 1) * self.r]
+                        for d in range(res.shape[1])
+                    ),
+                )
+                for m in range(self.m)
+            ],
+            axis=-1,
+        )  # compute each output dimension
+
+        # force (1,) output for non vectorial solution (consistency)
+        if len(res.shape) == self.d:
+            return jnp.expand_dims(res, axis=-1)
+        return res