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

jinns/parameters/_derivative_keys.py

@@ -2,73 +2,45 @@
 Formalize the data structure for the derivative keys
 """
 
-from functools import partial
-from dataclasses import fields, InitVar
+from dataclasses import InitVar
 from typing import Literal
+from jaxtyping import Array
 import jax
 import equinox as eqx
 
-from jinns.parameters._params import Params, ParamsDict
+from jinns.parameters._params import Params
 
 
 def _get_masked_parameters(
-    derivative_mask_str: str, params: Params | ParamsDict
-) -> Params | ParamsDict:
+    derivative_mask_str: str, params: Params[Array]
+) -> Params[bool]:
     """
     Creates the Params object with True values where we want to differentiate
     """
-    if isinstance(params, Params):
-        # start with a params object with True everywhere. We will update to False
-        # for parameters wrt which we do want not to differentiate the loss
-        diff_params = jax.tree.map(
-            lambda x: True,
-            params,
-            is_leaf=lambda x: isinstance(x, eqx.Module)
-            and not isinstance(x, Params),  # do not travers nn_params, more
-            # granularity could be imagined here, in the future
-        )
-        if derivative_mask_str == "both":
-            return diff_params
-        if derivative_mask_str == "eq_params":
-            return eqx.tree_at(lambda p: p.nn_params, diff_params, False)
-        if derivative_mask_str == "nn_params":
-            return eqx.tree_at(
-                lambda p: p.eq_params,
-                diff_params,
-                jax.tree.map(lambda x: False, params.eq_params),
-            )
-        raise ValueError(
-            "Bad value for DerivativeKeys. Got "
-            f'{derivative_mask_str}, expected "both", "nn_params" or '
-            ' "eq_params"'
-        )
-    elif isinstance(params, ParamsDict):
-        # do not travers nn_params, more
+    # start with a params object with True everywhere. We will update to False
+    # for parameters wrt which we do want not to differentiate the loss
+    diff_params = jax.tree.map(
+        lambda x: True,
+        params,
+        is_leaf=lambda x: isinstance(x, eqx.Module)
+        and not isinstance(x, Params),  # do not travers nn_params, more
         # granularity could be imagined here, in the future
-        diff_params = ParamsDict(
-            nn_params=True, eq_params=jax.tree.map(lambda x: True, params.eq_params)
-        )
-        if derivative_mask_str == "both":
-            return diff_params
-        if derivative_mask_str == "eq_params":
-            return eqx.tree_at(lambda p: p.nn_params, diff_params, False)
-        if derivative_mask_str == "nn_params":
-            return eqx.tree_at(
-                lambda p: p.eq_params,
-                diff_params,
-                jax.tree.map(lambda x: False, params.eq_params),
-            )
-        raise ValueError(
-            "Bad value for DerivativeKeys. Got "
-            f'{derivative_mask_str}, expected "both", "nn_params" or '
-            ' "eq_params"'
-        )
-
-    else:
-        raise ValueError(
-            f"Bad value for params. Got {type(params)}, expected Params "
-            " or ParamsDict"
+    )
+    if derivative_mask_str == "both":
+        return diff_params
+    if derivative_mask_str == "eq_params":
+        return eqx.tree_at(lambda p: p.nn_params, diff_params, False)
+    if derivative_mask_str == "nn_params":
+        return eqx.tree_at(
+            lambda p: p.eq_params,
+            diff_params,
+            jax.tree.map(lambda x: False, params.eq_params),
         )
+    raise ValueError(
+        "Bad value for DerivativeKeys. Got "
+        f'{derivative_mask_str}, expected "both", "nn_params" or '
+        ' "eq_params"'
+    )
 
 
 class DerivativeKeysODE(eqx.Module):
@@ -89,7 +61,7 @@ class DerivativeKeysODE(eqx.Module):
          1. For unspecified loss term, the default is to differentiate with
         respect to `"nn_params"` only.
          2. No granularity inside `Params.nn_params` is currently supported.
-         3. Note that the main Params or ParamsDict object of the problem is mandatory if initialization via `from_str()`.
+         3. Note that the main Params object of the problem is mandatory if initialization via `from_str()`.
 
     A typical specification is of the form:
     ```python
@@ -105,67 +77,69 @@ class DerivativeKeysODE(eqx.Module):
 
     Parameters
     ----------
-    dyn_loss : Params | ParamsDict | None, default=None
+    dyn_loss : Params[bool] | None, default=None
         Tell wrt which node of `Params` we will differentiate the
         dynamic loss. To do so, the fields of `Params` contain True (if
         differentiation) or False (if no differentiation).
-    observations : Params | ParamsDict | None, default=None
+    observations : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         observation loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    initial_condition : Params | ParamsDict | None, default=None
+    initial_condition : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         initial condition loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    params : InitVar[Params | ParamsDict], default=None
+    params : InitVar[Params[Array]], default=None
         The main Params object of the problem. It is required
         if some terms are unspecified (None). This is because, jinns cannot
         infer the content of `Params.eq_params`.
     """
 
-    dyn_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
-    observations: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
-    initial_condition: Params | ParamsDict | None = eqx.field(
-        kw_only=True, default=None
-    )
-
-    params: InitVar[Params | ParamsDict] = eqx.field(kw_only=True, default=None)
-
-    def __post_init__(self, params=None):
+    dyn_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    observations: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    initial_condition: Params[bool] | None = eqx.field(kw_only=True, default=None)
+
+    params: InitVar[Params[Array] | None] = eqx.field(kw_only=True, default=None)
+
+    def __post_init__(self, params: Params[Array] | None = None):
+        if params is None and (
+            self.dyn_loss is None
+            or self.observations is None
+            or self.initial_condition is None
+        ):
+            raise ValueError(
+                "params cannot be None since at least one loss "
+                "term has an undefined derivative key Params PyTree"
+            )
         if self.dyn_loss is None:
-            try:
-                self.dyn_loss = _get_masked_parameters("nn_params", params)
-            except AttributeError:
-                raise ValueError(
-                    "self.dyn_loss is None, hence params should be " "passed"
-                )
+            if params is None:
+                raise ValueError("self.dyn_loss is None, hence params should be passed")
+            self.dyn_loss = _get_masked_parameters("nn_params", params)
         if self.observations is None:
-            try:
-                self.observations = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
-                    "self.observations is None, hence params should be " "passed"
+                    "self.observations is None, hence params should be passed"
                 )
+            self.observations = _get_masked_parameters("nn_params", params)
         if self.initial_condition is None:
-            try:
-                self.initial_condition = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
-                    "self.initial_condition is None, hence params should be " "passed"
+                    "self.initial_condition is None, hence params should be passed"
                 )
+            self.initial_condition = _get_masked_parameters("nn_params", params)
 
     @classmethod
     def from_str(
         cls,
-        params: Params | ParamsDict,
+        params: Params[Array],
         dyn_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         observations: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         initial_condition: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
     ):
         """
@@ -181,19 +155,19 @@ class DerivativeKeysODE(eqx.Module):
         Parameters
         ----------
         params
-            The actual Params or ParamsDict object of the problem.
+            The actual Params object of the problem.
         dyn_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the dynamic loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         observations
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the observations. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         initial_condition
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the initial condition. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         """
         return DerivativeKeysODE(
             dyn_loss=(
@@ -220,78 +194,74 @@ class DerivativeKeysPDEStatio(eqx.Module):
 
     Parameters
     ----------
-     dyn_loss : Params | ParamsDict | None, default=None
+     dyn_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         dynamic loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-     observations : Params | ParamsDict | None, default=None
+     observations : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         observation loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-     boundary_loss : Params | ParamsDict | None, default=None
+     boundary_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         boundary loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-     norm_loss : Params | ParamsDict | None, default=None
+     norm_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         normalization loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-     params : InitVar[Params | ParamsDict], default=None
+     params : InitVar[Params[Array]], default=None
         The main Params object of the problem. It is required
         if some terms are unspecified (None). This is because, jinns cannot infer the
         content of `Params.eq_params`.
     """
 
-    dyn_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
-    observations: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
-    boundary_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
-    norm_loss: Params | ParamsDict | None = eqx.field(kw_only=True, default=None)
+    dyn_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    observations: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    boundary_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
+    norm_loss: Params[bool] | None = eqx.field(kw_only=True, default=None)
 
-    params: InitVar[Params | ParamsDict] = eqx.field(kw_only=True, default=None)
+    params: InitVar[Params[Array] | None] = eqx.field(kw_only=True, default=None)
 
-    def __post_init__(self, params=None):
+    def __post_init__(self, params: Params[Array] | None = None):
         if self.dyn_loss is None:
-            try:
-                self.dyn_loss = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError("self.dyn_loss is None, hence params should be passed")
+            self.dyn_loss = _get_masked_parameters("nn_params", params)
         if self.observations is None:
-            try:
-                self.observations = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
                     "self.observations is None, hence params should be passed"
                 )
+            self.observations = _get_masked_parameters("nn_params", params)
         if self.boundary_loss is None:
-            try:
-                self.boundary_loss = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
                     "self.boundary_loss is None, hence params should be passed"
                 )
+            self.boundary_loss = _get_masked_parameters("nn_params", params)
         if self.norm_loss is None:
-            try:
-                self.norm_loss = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
                     "self.norm_loss is None, hence params should be passed"
                 )
+            self.norm_loss = _get_masked_parameters("nn_params", params)
 
     @classmethod
     def from_str(
         cls,
-        params: Params | ParamsDict,
+        params: Params[Array],
         dyn_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         observations: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         boundary_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         norm_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
     ):
         """
@@ -300,23 +270,23 @@ class DerivativeKeysPDEStatio(eqx.Module):
         Parameters
         ----------
         params
-            The actual Param or ParamsDict object of the problem.
+            The actual Param object of the problem.
         dyn_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the dynamic loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         observations
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the observations. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         boundary_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the boundary loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         norm_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the normalization loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         """
         return DerivativeKeysPDEStatio(
             dyn_loss=(
@@ -348,64 +318,61 @@ class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
 
     Parameters
     ----------
-    dyn_loss : Params | ParamsDict | None, default=None
+    dyn_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         dynamic loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    observations : Params | ParamsDict | None, default=None
+    observations : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         observation loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    boundary_loss : Params | ParamsDict | None, default=None
+    boundary_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         boundary loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    norm_loss : Params | ParamsDict | None, default=None
+    norm_loss : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         normalization loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    initial_condition : Params | ParamsDict | None, default=None
+    initial_condition : Params[bool] | None, default=None
         Tell wrt which parameters among Params we will differentiate the
         initial_condition loss. To do so, the fields of Params contain True (if
         differentiation) or False (if no differentiation).
-    params : InitVar[Params | ParamsDict], default=None
+    params : InitVar[Params[Array]], default=None
         The main Params object of the problem. It is required
         if some terms are unspecified (None). This is because, jinns cannot infer the
         content of `Params.eq_params`.
     """
 
-    initial_condition: Params | ParamsDict | None = eqx.field(
-        kw_only=True, default=None
-    )
+    initial_condition: Params[bool] | None = eqx.field(kw_only=True, default=None)
 
-    def __post_init__(self, params=None):
+    def __post_init__(self, params: Params[Array] | None = None):
         super().__post_init__(params=params)
         if self.initial_condition is None:
-            try:
-                self.initial_condition = _get_masked_parameters("nn_params", params)
-            except AttributeError:
+            if params is None:
                 raise ValueError(
                     "self.initial_condition is None, hence params should be passed"
                 )
+            self.initial_condition = _get_masked_parameters("nn_params", params)
 
     @classmethod
     def from_str(
         cls,
-        params: Params | ParamsDict,
+        params: Params[Array],
         dyn_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         observations: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         boundary_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         norm_loss: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
         initial_condition: (
-            Literal["nn_params", "eq_params", "both"] | Params | ParamsDict
+            Literal["nn_params", "eq_params", "both"] | Params[bool]
         ) = "nn_params",
     ):
         """
@@ -414,27 +381,27 @@ class DerivativeKeysPDENonStatio(DerivativeKeysPDEStatio):
         Parameters
         ----------
         params
-            The actual Params | ParamsDict object of the problem.
+            The actual Params object of the problem.
         dyn_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the dynamic loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         observations
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the observations. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         boundary_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the boundary loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         norm_loss
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the normalization loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         initial_condition
             Tell wrt which parameters among `"nn_params"`, `"eq_params"` or
             `"both"` we will differentiate the initial_condition loss. Default is
-            `"nn_params"`. Specifying a Params or ParamsDict is also possible.
+            `"nn_params"`. Specifying a Params is also possible.
         """
         return DerivativeKeysPDENonStatio(
             dyn_loss=(
@@ -471,32 +438,6 @@ def _set_derivatives(params, derivative_keys):
     has a copy of the params with appropriate derivatives set
     """
 
-    def _set_derivatives_ParamsDict(params_, derivative_mask):
-        """
-        The next lines put a stop_gradient around the fields that do not
-        differentiate the loss term
-        **Note:** **No granularity inside `ParamsDict.nn_params` is currently
-        supported.**
-        This means a typical Params specification is of the form:
-        `ParamsDict(nn_params=True | False, eq_params={"0":{"alpha":True | False,
-        "beta":True | False}}, "1":{"alpha":True | False, "beta":True | False}})`.
-        """
-        # a ParamsDict object is reconstructed by hand since we do not want to
-        # traverse nn_params, for now...
-        return ParamsDict(
-            nn_params=jax.lax.cond(
-                derivative_mask.nn_params,
-                lambda p: p,
-                jax.lax.stop_gradient,
-                params_.nn_params,
-            ),
-            eq_params=jax.tree.map(
-                lambda p, d: jax.lax.cond(d, lambda p: p, jax.lax.stop_gradient, p),
-                params_.eq_params,
-                derivative_mask.eq_params,
-            ),
-        )
-
     def _set_derivatives_(params_, derivative_mask):
         """
         The next lines put a stop_gradient around the fields that do not
@@ -516,6 +457,4 @@ def _set_derivatives(params, derivative_keys):
             # granularity could be imagined here, in the future
         )
 
-    if isinstance(params, ParamsDict):
-        return _set_derivatives_ParamsDict(params, derivative_keys)
     return _set_derivatives_(params, derivative_keys)

jinns/parameters/_params.py

@@ -2,67 +2,36 @@
 Formalize the data structure for the parameters
 """
 
+from typing import Generic, TypeVar
 import jax
 import equinox as eqx
-from typing import Dict
-from jaxtyping import Array, PyTree
+from jaxtyping import Array, PyTree, Float
 
+T = TypeVar("T")  # the generic type for what is in the Params PyTree because we
+# have possibly Params of Arrays, boolean, ...
 
-class Params(eqx.Module):
+
+class Params(eqx.Module, Generic[T]):
     """
     The equinox module for the parameters
 
     Parameters
     ----------
-    nn_params : Pytree
+    nn_params : PyTree[T]
         A PyTree of the non-static part of the PINN eqx.Module, i.e., the
         parameters of the PINN
-    eq_params : Dict[str, Array]
+    eq_params : dict[str, T]
         A dictionary of the equation parameters. Keys are the parameter name,
         values are their corresponding value
     """
 
-    nn_params: PyTree = eqx.field(kw_only=True, default=None)
-    eq_params: Dict[str, Array] = eqx.field(kw_only=True, default=None)
-
-
-class ParamsDict(eqx.Module):
-    """
-    The equinox module for a dictionnary of parameters with different keys
-    corresponding to different equations.
-
-    Parameters
-    ----------
-    nn_params : Dict[str, PyTree]
-        The neural network's parameters. Most of the time, it will be the
-        Array part of an `eqx.Module` obtained by
-        `eqx.partition(module, eqx.is_inexact_array)`.
-    eq_params : Dict[str, Array]
-        A dictionary of the equation parameters. Dict keys are the parameter name as defined your custom loss.
-    """
-
-    nn_params: Dict[str, PyTree] = eqx.field(kw_only=True, default=None)
-    eq_params: Dict[str, Array] = eqx.field(kw_only=True, default=None)
-
-    def extract_params(self, nn_key: str) -> Params:
-        """
-        Extract the corresponding `nn_params` and `eq_params` for `nn_key` and
-        return them in the form of a `Params` object.
-        """
-        try:
-            return Params(
-                nn_params=self.nn_params[nn_key],
-                eq_params=self.eq_params[nn_key],
-            )
-        except (KeyError, IndexError) as e:
-            return Params(
-                nn_params=self.nn_params[nn_key],
-                eq_params=self.eq_params,
-            )
+    nn_params: PyTree[T] = eqx.field(kw_only=True, default=None)
+    eq_params: dict[str, T] = eqx.field(kw_only=True, default=None)
 
 
 def _update_eq_params_dict(
-    params: Params, param_batch_dict: Dict[str, Array]
+    params: Params[Array],
+    param_batch_dict: dict[str, Float[Array, " param_batch_size dim"]],
 ) -> Params:
     """
     Update params.eq_params with a batch of eq_params for given key(s)
@@ -89,13 +58,16 @@ def _update_eq_params_dict(
 
 
 def _get_vmap_in_axes_params(
-    eq_params_batch_dict: Dict[str, Array], params: Params | ParamsDict
-) -> tuple[Params]:
+    eq_params_batch_dict: dict[str, Array], params: Params[Array]
+) -> tuple[Params[int | None] | None]:
     """
     Return the input vmap axes when there is batch(es) of parameters to vmap
     over. The latter are designated by keys in eq_params_batch_dict.
     If eq_params_batch_dict is None (i.e. no additional parameter batch), we
     return (None,).
+
+    Note that we return a Params PyTree with an integer to designate the
+    vmapped axis or None if there is not
     """
     if eq_params_batch_dict is None:
         return (None,)
@@ -104,7 +76,7 @@ def _get_vmap_in_axes_params(
     # this is for a fine-grained vmaping
     # scheme over the params
     vmap_in_axes_params = (
-        type(params)(
+        Params(
             nn_params=None,
             eq_params={
                 k: (0 if k in eq_params_batch_dict.keys() else None)

jinns/plot/__init__.py

@@ -3,3 +3,5 @@ from jinns.plot._plot import (
     plot1d_slice,
     plot1d_image,
 )
+
+__all__ = ["plot2d", "plot1d_slice", "plot1d_image"]