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

jinns/loss/_operators.py

@@ -2,22 +2,42 @@
 Implements diverse operators for dynamic losses
 """
 
-from typing import Literal
+from __future__ import (
+    annotations,
+)
+
+from typing import Literal, cast, Callable
 
 import jax
 import jax.numpy as jnp
 from jax import grad
-import equinox as eqx
 from jaxtyping import Float, Array
 from jinns.parameters._params import Params
+from jinns.nn._abstract_pinn import AbstractPINN
+
+
+def _get_eq_type(
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None,
+) -> Literal["nonstatio_PDE", "statio_PDE"]:
+    """
+    But we filter out ODE from eq_type because we only have operators that does
+    not work with ODEs so far
+    """
+    if isinstance(u, AbstractPINN):
+        assert u.eq_type != "ODE", "Cannot compute the operator for ODE PINNs"
+        return u.eq_type
+    if eq_type is None:
+        raise ValueError("eq_type could not be set!")
+    return eq_type
 
 
 def divergence_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> float:
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " "]:
     r"""
     Compute the divergence of a vector field $\mathbf{u}$, i.e.,
     $\nabla_\mathbf{x} \cdot \mathbf{u}(\mathrm{inputs})$ with $\mathbf{u}$ a vector
@@ -41,13 +61,7 @@ def divergence_rev(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
         if eq_type == "nonstatio_PDE":
@@ -70,11 +84,11 @@ def divergence_rev(
 
 
 def divergence_fwd(
-    inputs: Float[Array, "batch_size dim"] | Float[Array, "batch_size 1+dim"],
-    u: eqx.Module,
-    params: Params,
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> Float[Array, "batch_size * (1+dim) 1"] | Float[Array, "batch_size * (dim) 1"]:
+    inputs: Float[Array, " batch_size dim"] | Float[Array, " batch_size 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the divergence of a **batched** vector field $\mathbf{u}$, i.e.,
     $\nabla_\mathbf{x} \cdot \mathbf{u}(\mathbf{x})$ with $\mathbf{u}$ a vector
@@ -103,13 +117,7 @@ def divergence_fwd(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     def scan_fun(_, i):
         if eq_type == "nonstatio_PDE":
@@ -142,12 +150,12 @@ def divergence_fwd(
 
 
 def laplacian_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
     method: Literal["trace_hessian_x", "trace_hessian_t_x", "loop"] = "trace_hessian_x",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> float:
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " "]:
     r"""
     Compute the Laplacian of a scalar field $u$ from $\mathbb{R}^d$
     to $\mathbb{R}$ or from $\mathbb{R}^{1+d}$ to $\mathbb{R}$, i.e., this
@@ -180,13 +188,7 @@ def laplacian_rev(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     if method == "trace_hessian_x":
         # NOTE we afford a concatenate here to avoid computing Hessian elements for
@@ -226,16 +228,12 @@ def laplacian_rev(
             if eq_type == "nonstatio_PDE":
                 d2u_dxi2 = grad(
                     lambda inputs: grad(u_)(inputs)[1 + i],
-                )(
-                    inputs
-                )[1 + i]
+                )(inputs)[1 + i]
             else:
                 d2u_dxi2 = grad(
                     lambda inputs: grad(u_, 0)(inputs)[i],
                     0,
-                )(
-                    inputs
-                )[i]
+                )(inputs)[i]
             return _, d2u_dxi2
 
         if eq_type == "nonstatio_PDE":
@@ -251,12 +249,12 @@ def laplacian_rev(
 
 
 def laplacian_fwd(
-    inputs: Float[Array, "batch_size 1+dim"] | Float[Array, "batch_size dim"],
-    u: eqx.Module,
-    params: Params,
+    inputs: Float[Array, " batch_size 1+dim"] | Float[Array, " batch_size dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
     method: Literal["trace_hessian_t_x", "trace_hessian_x", "loop"] = "loop",
-    eq_type: Literal["nonstatio_PDE", "statio_PDE"] = None,
-) -> Float[Array, "batch_size * (1+dim) 1"] | Float[Array, "batch_size * (dim) 1"]:
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) 1"] | Float[Array, " batch_size * (dim) 1"]:
     r"""
     Compute the Laplacian of a **batched** scalar field $u$
     from $\mathbb{R}^{b\times d}$ to $\mathbb{R}^{b\times b}$ or
@@ -299,13 +297,7 @@ def laplacian_fwd(
         can know that by inspecting the `u` argument (PINN object). But if `u` is
         a function, we must set this attribute.
     """
-
-    try:
-        eq_type = u.eq_type
-    except AttributeError:
-        pass  # use the value passed as argument
-    if eq_type is None:
-        raise ValueError("eq_type could not be set!")
+    eq_type = _get_eq_type(u, eq_type)
 
     if method == "loop":
 
@@ -398,11 +390,12 @@ def laplacian_fwd(
 
 
 def vectorial_laplacian_rev(
-    inputs: Float[Array, "dim"] | Float[Array, "1+dim"],
-    u: eqx.Module,
-    params: Params,
-    dim_out: int = None,
-) -> Float[Array, "dim_out"]:
+    inputs: Float[Array, " dim"] | Float[Array, " 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    dim_out: int | None = None,
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " dim_out"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ from
     $\mathbb{R}^d$ to $\mathbb{R}^n$ or from $\mathbb{R}^{1+d}$ to
@@ -426,7 +419,12 @@ def vectorial_laplacian_rev(
     dim_out
         Dimension of the vector $\mathbf{u}(\mathrm{inputs})$. This needs to be
         provided if it is different than that of $\mathrm{inputs}$.
+    eq_type
+        whether we consider a stationary or non stationary PINN. Most often we
+        can know that by inspecting the `u` argument (PINN object). But if `u` is
+        a function, we must set this attribute.
     """
+    eq_type = _get_eq_type(u, eq_type)
     if dim_out is None:
         dim_out = inputs.shape[0]
 
@@ -435,7 +433,9 @@ def vectorial_laplacian_rev(
         # each of these components
         # Note the jnp.expand_dims call
         uj = lambda inputs, params: jnp.expand_dims(u(inputs, params)[j], axis=-1)
-        lap_on_j = laplacian_rev(inputs, uj, params, eq_type=u.eq_type)
+        lap_on_j = laplacian_rev(
+            inputs, cast(AbstractPINN, uj), params, eq_type=eq_type
+        )
 
         return _, lap_on_j
 
@@ -444,11 +444,12 @@ def vectorial_laplacian_rev(
 
 
 def vectorial_laplacian_fwd(
-    inputs: Float[Array, "batch_size dim"] | Float[Array, "batch_size 1+dim"],
-    u: eqx.Module,
-    params: Params,
-    dim_out: int = None,
-) -> Float[Array, "batch_size * (1+dim) n"] | Float[Array, "batch_size * (dim) n"]:
+    inputs: Float[Array, " batch_size dim"] | Float[Array, " batch_size 1+dim"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+    dim_out: int | None = None,
+    eq_type: Literal["nonstatio_PDE", "statio_PDE"] | None = None,
+) -> Float[Array, " batch_size * (1+dim) n"] | Float[Array, " batch_size * (dim) n"]:
     r"""
     Compute the vectorial Laplacian of a vector field $\mathbf{u}$ when
     `u` is a SPINN, in this case, it corresponds to a vector
@@ -474,7 +475,12 @@ def vectorial_laplacian_fwd(
     dim_out
         the value of the output dimension ($n$ in the formula above). Must be
         set if different from $d$.
+    eq_type
+        whether we consider a stationary or non stationary PINN. Most often we
+        can know that by inspecting the `u` argument (PINN object). But if `u` is
+        a function, we must set this attribute.
     """
+    eq_type = _get_eq_type(u, eq_type)
     if dim_out is None:
         dim_out = inputs.shape[0]
 
@@ -483,7 +489,9 @@ def vectorial_laplacian_fwd(
         # each of these components
         # Note the expand_dims
         uj = lambda inputs, params: jnp.expand_dims(u(inputs, params)[..., j], axis=-1)
-        lap_on_j = laplacian_fwd(inputs, uj, params, eq_type=u.eq_type)
+        lap_on_j = laplacian_fwd(
+            inputs, cast(AbstractPINN, uj), params, eq_type=eq_type
+        )
 
         return _, lap_on_j
 
@@ -492,8 +500,10 @@ def vectorial_laplacian_fwd(
 
 
 def _u_dot_nabla_times_u_rev(
-    x: Float[Array, "2"], u: eqx.Module, params: Params
-) -> Float[Array, "2"]:
+    x: Float[Array, " 2"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+) -> Float[Array, " 2"]:
     r"""
     Implement $((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})$ for
     $\mathbf{x}$ of arbitrary
@@ -522,10 +532,10 @@ def _u_dot_nabla_times_u_rev(
 
 
 def _u_dot_nabla_times_u_fwd(
-    x: Float[Array, "batch_size 2"],
-    u: eqx.Module,
-    params: Params,
-) -> Float[Array, "batch_size batch_size 2"]:
+    x: Float[Array, " batch_size 2"],
+    u: AbstractPINN | Callable[[Array, Params[Array]], Array],
+    params: Params[Array],
+) -> Float[Array, " batch_size batch_size 2"]:
     r"""
     Implement :math:`((\mathbf{u}\cdot\nabla)\mathbf{u})(\mathbf{x})` for
     :math:`\mathbf{x}` of arbitrary dimension **with a batch dimension**.

jinns/nn/__init__.py

@@ -1,7 +1,22 @@
 from ._save_load import save_pinn, load_pinn
+from ._abstract_pinn import AbstractPINN
 from ._pinn import PINN
 from ._spinn import SPINN
 from ._mlp import PINN_MLP, MLP
 from ._spinn_mlp import SPINN_MLP, SMLP
 from ._hyperpinn import HyperPINN
 from ._ppinn import PPINN_MLP
+
+__all__ = [
+    "save_pinn",
+    "load_pinn",
+    "AbstractPINN",
+    "PINN",
+    "SPINN",
+    "PINN_MLP",
+    "MLP",
+    "SPINN_MLP",
+    "SMLP",
+    "HyperPINN",
+    "PPINN_MLP",
+]

jinns/nn/_abstract_pinn.py

@@ -0,0 +1,22 @@
+import abc
+from typing import Literal, Any
+from jaxtyping import Array
+import equinox as eqx
+
+from jinns.nn._utils import _PyTree_to_Params
+from jinns.parameters._params import Params
+
+
+class AbstractPINN(eqx.Module):
+    """
+    Basically just a way to add a __call__ to an eqx.Module.
+    The way to go for correct type hints apparently
+    https://github.com/patrick-kidger/equinox/issues/1002 + https://docs.kidger.site/equinox/pattern/
+    """
+
+    eq_type: eqx.AbstractVar[Literal["ODE", "statio_PDE", "nonstatio_PDE"]]
+
+    @abc.abstractmethod
+    @_PyTree_to_Params
+    def __call__(self, inputs: Any, params: Params[Array], *args, **kwargs) -> Any:
+        pass

jinns/nn/_hyperpinn.py

@@ -3,9 +3,11 @@ Implements utility function to create HyperPINNs
 https://arxiv.org/pdf/2111.01008.pdf
 """
 
+from __future__ import annotations
+
 import warnings
 from dataclasses import InitVar
-from typing import Callable, Literal, Self, Union, Any
+from typing import Callable, Literal, Self, Union, Any, cast, overload
 from math import prod
 import jax
 import jax.numpy as jnp
@@ -15,12 +17,13 @@ import numpy as onp
 
 from jinns.nn._pinn import PINN
 from jinns.nn._mlp import MLP
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
+from jinns.nn._utils import _PyTree_to_Params
 
 
 def _get_param_nb(
-    params: Params,
-) -> tuple[int, list]:
+    params: PyTree[Array],
+) -> tuple[int, list[int]]:
     """Returns the number of parameters in a Params object and also
     the cumulative sum when parsing the object.
 
@@ -48,7 +51,7 @@ class HyperPINN(PINN):
 
     Parameters
     ----------
-    hyperparams: list = eqx.field(static=True)
+    hyperparams: list[str] = eqx.field(static=True)
         A list of keys from Params.eq_params that will be considered as
         hyperparameters for metamodeling.
     hypernet_input_size: int
@@ -72,12 +75,12 @@ class HyperPINN(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 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.
@@ -100,10 +103,10 @@ class HyperPINN(PINN):
     eqx_hyper_network: InitVar[eqx.Module] = eqx.field(kw_only=True)
 
     pinn_params_sum: int = eqx.field(init=False, static=True)
-    pinn_params_cumsum: list = eqx.field(init=False, static=True)
+    pinn_params_cumsum: list[int] = eqx.field(init=False, static=True)
 
-    init_params_hyper: PyTree = eqx.field(init=False)
-    static_hyper: PyTree = eqx.field(init=False, static=True)
+    init_params_hyper: HyperPINN = eqx.field(init=False)
+    static_hyper: HyperPINN = eqx.field(init=False, static=True)
 
     def __post_init__(self, eqx_network, eqx_hyper_network):
         super().__post_init__(
@@ -115,7 +118,7 @@ class HyperPINN(PINN):
         )
         self.pinn_params_sum, self.pinn_params_cumsum = _get_param_nb(self.init_params)
 
-    def _hyper_to_pinn(self, hyper_output: Float[Array, "output_dim"]) -> PyTree:
+    def _hyper_to_pinn(self, hyper_output: Float[Array, "  output_dim"]) -> PINN:
         """
         From the output of the hypernetwork, transform to a well formed
         parameters for the pinn network (i.e. with the same PyTree structure as
@@ -142,15 +145,29 @@ class HyperPINN(PINN):
             is_leaf=lambda x: isinstance(x, jnp.ndarray),
         )
 
+    @overload
+    @_PyTree_to_Params
     def __call__(
         self,
-        inputs: Float[Array, "input_dim"],
-        params: Params | ParamsDict | PyTree,
+        inputs: Float[Array, " input_dim"],
+        params: PyTree,
         *args,
         **kwargs,
-    ) -> Float[Array, "output_dim"]:
+    ) -> Float[Array, " output_dim"]: ...
+
+    @_PyTree_to_Params
+    def __call__(
+        self,
+        inputs: Float[Array, "  input_dim"],
+        params: Params[Array],
+        *args,
+        **kwargs,
+    ) -> Float[Array, "  output_dim"]:
         """
         Evaluate the HyperPINN 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
@@ -158,16 +175,17 @@ class HyperPINN(PINN):
             # DataGenerators)
             inputs = inputs[None]
 
-        try:
-            hyper = eqx.combine(params.nn_params, self.static_hyper)
-        except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
-            hyper = eqx.combine(params, self.static_hyper)
+        # try:
+        hyper = eqx.combine(params.nn_params, self.static_hyper)
+        # except (KeyError, AttributeError, TypeError) as e:  # give more flexibility
+        #    hyper = eqx.combine(params, self.static_hyper)
 
         eq_params_batch = jnp.concatenate(
-            [params.eq_params[k].flatten() for k in self.hyperparams], axis=0
+            [params.eq_params[k].flatten() for k in self.hyperparams],
+            axis=0,
         )
 
-        hyper_output = hyper(eq_params_batch)
+        hyper_output = hyper(eq_params_batch)  # type: ignore
 
         pinn_params = self._hyper_to_pinn(hyper_output)
 
@@ -187,21 +205,34 @@ class HyperPINN(PINN):
         eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
         hyperparams: list[str],
         hypernet_input_size: int,
-        eqx_network: eqx.nn.MLP = None,
-        eqx_hyper_network: eqx.nn.MLP = None,
+        eqx_network: eqx.nn.MLP | MLP | None = None,
+        eqx_hyper_network: eqx.nn.MLP | MLP | None = None,
         key: Key = None,
-        eqx_list: tuple[tuple[Callable, int, int] | Callable, ...] = None,
-        eqx_list_hyper: 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,
+        eqx_list_hyper: (
+            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, HyperPINN]:
         r"""
         Utility function to create a standard PINN neural network with the equinox
         library.
@@ -250,11 +281,11 @@ class HyperPINN(PINN):
             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)
             )`.
         eqx_list_hyper
@@ -268,11 +299,11 @@ class HyperPINN(PINN):
             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)
             )`.
         input_transform
@@ -343,10 +374,13 @@ class HyperPINN(PINN):
                     (eqx_list_hyper[-1][:2] + (pinn_params_sum,)),
                 )
             else:
-                eqx_list_hyper = (
-                    eqx_list_hyper[:-2]
-                    + ((eqx_list_hyper[-2][:2] + (pinn_params_sum,)),)
-                    + eqx_list_hyper[-1]
+                eqx_list_hyper = cast(
+                    tuple[tuple[Callable, int, int] | tuple[Callable], ...],
+                    (
+                        eqx_list_hyper[:-2]
+                        + ((eqx_list_hyper[-2][:2] + (pinn_params_sum,)),)
+                        + eqx_list_hyper[-1]
+                    ),
                 )
             if len(eqx_list_hyper[0]) > 1:
                 eqx_list_hyper = (
@@ -357,21 +391,24 @@ class HyperPINN(PINN):
                     ),
                 ) + eqx_list_hyper[1:]
             else:
-                eqx_list_hyper = (
-                    eqx_list_hyper[0]
-                    + (
-                        (
-                            (eqx_list_hyper[1][0],)
-                            + (hypernet_input_size,)
-                            + (eqx_list_hyper[1][2],)
-                        ),
-                    )
-                    + eqx_list_hyper[2:]
+                eqx_list_hyper = cast(
+                    tuple[tuple[Callable, int, int] | tuple[Callable], ...],
+                    (
+                        eqx_list_hyper[0]
+                        + (
+                            (
+                                (eqx_list_hyper[1][0],)
+                                + (hypernet_input_size,)
+                                + (eqx_list_hyper[1][2],)  # type: ignore because we suppose that the second element of tuple is nec.of length > 1 since we expect smth like eqx.nn.Linear
+                            ),
+                        )
+                        + eqx_list_hyper[2:]
+                    ),
                 )
             key, subkey = jax.random.split(key, 2)
             # with warnings.catch_warnings():
             #    warnings.filterwarnings("ignore", message="A JAX array is being set as static!")
-            eqx_hyper_network = MLP(key=subkey, eqx_list=eqx_list_hyper)
+            eqx_hyper_network = cast(MLP, MLP(key=subkey, eqx_list=eqx_list_hyper))
 
             ### End of finetuning the hypernetwork architecture
 
@@ -386,10 +423,10 @@ class HyperPINN(PINN):
             hyperpinn = cls(
                 eqx_network=eqx_network,
                 eqx_hyper_network=eqx_hyper_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
                 hyperparams=hyperparams,
                 hypernet_input_size=hypernet_input_size,
                 filter_spec=filter_spec,