jinns 1.5.1__py3-none-any.whl → 1.6.1__py3-none-any.whl

jinns/data/_DataGeneratorODE.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Key, Array, Float
+from jaxtyping import PRNGKeyArray, Array, Float
 from jinns.data._Batchs import ODEBatch
 from jinns.data._utils import _check_and_set_rar_parameters, _reset_or_increment
 from jinns.data._AbstractDataGenerator import AbstractDataGenerator
@@ -24,7 +24,7 @@ class DataGeneratorODE(AbstractDataGenerator):
 
     Parameters
     ----------
-    key : Key
+    key : PRNGKeyArray
         Jax random key to sample new time points and to shuffle batches
     nt : int
         The number of total time points that will be divided in
@@ -42,10 +42,10 @@ class DataGeneratorODE(AbstractDataGenerator):
         The method that generates the `nt` time points. `grid` means
         regularly spaced points over the domain. `uniform` means uniformly
         sampled points over the domain
-    rar_parameters : RarParameterDict, default=None
+    rar_parameters : None | RarParameterDict, default=None
        A TypedDict to specify the Residual Adaptative Resampling procedure. See
        the docstring from RarParameterDict
-    n_start : int, default=None
+    n_start : None | int, default=None
         Defaults to None. The effective size of nt used at start time.
         This value must be
         provided when rar_parameters is not None. Otherwise we set internally
@@ -54,25 +54,43 @@ class DataGeneratorODE(AbstractDataGenerator):
         then corresponds to the initial number of points we train the PINN.
     """
 
-    key: Key = eqx.field(kw_only=True)
-    nt: int = eqx.field(kw_only=True, static=True)
-    tmin: Float = eqx.field(kw_only=True)
-    tmax: Float = eqx.field(kw_only=True)
-    temporal_batch_size: int | None = eqx.field(static=True, default=None, kw_only=True)
-    method: str = eqx.field(
-        static=True, kw_only=True, default_factory=lambda: "uniform"
-    )
-    rar_parameters: dict[str, int] = eqx.field(default=None, kw_only=True)
-    n_start: int = eqx.field(static=True, default=None, kw_only=True)
-
-    # all the init=False fields are set in __post_init__
+    key: PRNGKeyArray
+    nt: int = eqx.field(static=True)
+    tmin: float
+    tmax: float
+    temporal_batch_size: int | None = eqx.field(static=True)
+    method: str = eqx.field(static=True)
+    rar_parameters: None | dict[str, int]
+    n_start: None | int
+
+    # --- Below fields are not passed as arguments to __init__
     p: Float[Array, " nt 1"] | None = eqx.field(init=False)
     rar_iter_from_last_sampling: int | None = eqx.field(init=False)
     rar_iter_nb: int | None = eqx.field(init=False)
     curr_time_idx: int = eqx.field(init=False)
     times: Float[Array, " nt 1"] = eqx.field(init=False)
 
-    def __post_init__(self):
+    def __init__(
+        self,
+        *,
+        key: PRNGKeyArray,
+        nt: int,
+        tmin: float,
+        tmax: float,
+        temporal_batch_size: int | None,
+        method: str = "uniform",
+        rar_parameters: None | dict[str, int] = None,
+        n_start: None | int = None,
+    ):
+        self.key = key
+        self.nt = nt
+        self.tmin = tmin
+        self.tmax = tmax
+        self.temporal_batch_size = temporal_batch_size
+        self.method = method
+        self.n_start = n_start
+        self.rar_parameters = rar_parameters
+
         (
             self.n_start,
             self.p,
@@ -97,7 +115,7 @@ class DataGeneratorODE(AbstractDataGenerator):
         # above way for the key.
 
     def sample_in_time_domain(
-        self, key: Key, sample_size: int | None = None
+        self, key: PRNGKeyArray, sample_size: int | None = None
     ) -> Float[Array, " nt 1"]:
         return jax.random.uniform(
             key,
@@ -106,7 +124,9 @@ class DataGeneratorODE(AbstractDataGenerator):
             maxval=self.tmax,
         )
 
-    def generate_time_data(self, key: Key) -> tuple[Key, Float[Array, " nt"]]:
+    def generate_time_data(
+        self, key: PRNGKeyArray
+    ) -> tuple[PRNGKeyArray, Float[Array, " nt"]]:
         """
         Construct a complete set of `self.nt` time points according to the
         specified `self.method`
@@ -125,7 +145,11 @@ class DataGeneratorODE(AbstractDataGenerator):
     def _get_time_operands(
         self,
     ) -> tuple[
-        Key, Float[Array, " nt 1"], int, int | None, Float[Array, " nt 1"] | None
+        PRNGKeyArray,
+        Float[Array, " nt 1"],
+        int,
+        int | None,
+        Float[Array, " nt 1"] | None,
     ]:
         return (
             self.key,
@@ -150,7 +174,7 @@ class DataGeneratorODE(AbstractDataGenerator):
         bend = bstart + self.temporal_batch_size
 
         # Compute the effective number of used collocation points
-        if self.rar_parameters is not None:
+        if self.rar_parameters is not None and self.n_start is not None:
             nt_eff = (
                 self.n_start
                 + self.rar_iter_nb  # type: ignore
@@ -167,7 +191,9 @@ class DataGeneratorODE(AbstractDataGenerator):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.times, m.curr_time_idx), self, new_attributes
+            lambda m: (m.key, m.times, m.curr_time_idx),  # type: ignore
+            self,
+            new_attributes,
         )
 
         # commands below are equivalent to

jinns/data/_DataGeneratorObservations.py

@@ -8,10 +8,32 @@ from __future__ import (
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Key, Int, Array, Float
+from typing import TYPE_CHECKING, Self
+from jaxtyping import PRNGKeyArray, Int, Array, Float
 from jinns.data._Batchs import ObsBatchDict
 from jinns.data._utils import _reset_or_increment
 from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+from jinns.utils._DictToModuleMeta import DictToModuleMeta
+
+if TYPE_CHECKING:
+    # imports only used in type hints
+    InputEqParams = (
+        dict[str, Float[Array, "  n_obs"]] | dict[str, Float[Array, " n_obs 1"]]
+    ) | None
+
+    # Note that the lambda functions used below are with type: ignore just
+    # because the lambda are not type annotated, but there is no proper way
+    # to do this and we should assign the lambda to a type hinted variable
+    # before hand: this is not practical, let us not get mad at this
+
+
+class DGObservedParams(metaclass=DictToModuleMeta):
+    """
+    However, static type checkers cannot know that DGObservedParams inherit from
+    eqx.Module and explicit casting to the latter class will be needed
+    """
+
+    pass
 
 
 class DataGeneratorObservations(AbstractDataGenerator):
@@ -21,7 +43,7 @@ class DataGeneratorObservations(AbstractDataGenerator):
 
     Parameters
     ----------
-    key : Key
+    key : PRNGKeyArray
         Jax random key to shuffle batches
     obs_batch_size : int | None
         The size of the batch of randomly selected points among
@@ -39,7 +61,8 @@ class DataGeneratorObservations(AbstractDataGenerator):
     observed_eq_params : dict[str, Float[Array, " n_obs 1"]], default={}
         A dict with keys corresponding to
         the parameter name. The keys must match the keys in
-        `params["eq_params"]`. The values are jnp.array with 2 dimensions
+        `params["eq_params"]`, ie., if only some parameters are observed, other
+        keys **must still appear with None as value**. The values are jnp.array with 2 dimensions
         with values corresponding to the parameter value for which we also
         have observed_pinn_in and observed_values. Hence the first
         dimension must be aligned with observed_pinn_in and observed_values.
@@ -54,30 +77,37 @@ class DataGeneratorObservations(AbstractDataGenerator):
         arguments of `jinns.solve()`. Read `jinns.solve()` doc for more info.
     """
 
-    key: Key
+    key: PRNGKeyArray
     obs_batch_size: int | None = eqx.field(static=True)
     observed_pinn_in: Float[Array, " n_obs nb_pinn_in"]
     observed_values: Float[Array, " n_obs nb_pinn_out"]
-    observed_eq_params: dict[str, Float[Array, " n_obs 1"]] = eqx.field(
-        static=True, default_factory=lambda: {}
-    )
-    sharding_device: jax.sharding.Sharding = eqx.field(static=True, default=None)
+    observed_eq_params: eqx.Module | None
+    sharding_device: jax.sharding.Sharding | None  # = eqx.field(static=True)
 
     n: int = eqx.field(init=False, static=True)
     curr_idx: int = eqx.field(init=False)
     indices: Array = eqx.field(init=False)
 
-    def __post_init__(self):
+    def __init__(
+        self,
+        *,
+        key: PRNGKeyArray,
+        obs_batch_size: int | None = None,
+        observed_pinn_in: Float[Array, " n_obs nb_pinn_in"],
+        observed_values: Float[Array, " n_obs nb_pinn_out"],
+        observed_eq_params: InputEqParams | None = None,
+        sharding_device: jax.sharding.Sharding | None = None,
+    ) -> None:
+        super().__init__()
+        self.key = key
+        self.obs_batch_size = obs_batch_size
+        self.observed_pinn_in = observed_pinn_in
+        self.observed_values = observed_values
+
         if self.observed_pinn_in.shape[0] != self.observed_values.shape[0]:
             raise ValueError(
                 "self.observed_pinn_in and self.observed_values must have same first axis"
             )
-        for _, v in self.observed_eq_params.items():
-            if v.shape[0] != self.observed_pinn_in.shape[0]:
-                raise ValueError(
-                    "self.observed_pinn_in and the values of"
-                    " self.observed_eq_params must have the same first axis"
-                )
         if len(self.observed_pinn_in.shape) == 1:
             self.observed_pinn_in = self.observed_pinn_in[:, None]
         if self.observed_pinn_in.ndim > 2:
@@ -86,16 +116,34 @@ class DataGeneratorObservations(AbstractDataGenerator):
             self.observed_values = self.observed_values[:, None]
         if self.observed_values.ndim > 2:
             raise ValueError("self.observed_values must have 2 dimensions")
-        for k, v in self.observed_eq_params.items():
-            if len(v.shape) == 1:
-                self.observed_eq_params[k] = v[:, None]
-            if len(v.shape) > 2:
-                raise ValueError(
-                    "Each value of observed_eq_params must have 2 dimensions"
-                )
+
+        if observed_eq_params is not None:
+            for _, v in observed_eq_params.items():
+                if v.shape[0] != self.observed_pinn_in.shape[0]:
+                    raise ValueError(
+                        "self.observed_pinn_in and the values of"
+                        " self.observed_eq_params must have the same first axis"
+                    )
+            for k, v in observed_eq_params.items():
+                if len(v.shape) == 1:
+                    # Reshape to add an axis for 1-d Array
+                    observed_eq_params[k] = v[:, None]
+                if len(v.shape) > 2:
+                    raise ValueError(
+                        f"Each key of observed_eq_params must have 2"
+                        f"dimensions, key {k} had shape {v.shape}."
+                    )
+            # Convert the dict of observed parameters to the internal `EqParams`
+            # class used by Jinns.
+            self.observed_eq_params = DGObservedParams(
+                observed_eq_params, "DGObservedParams"
+            )
+        else:
+            self.observed_eq_params = observed_eq_params
 
         self.n = self.observed_pinn_in.shape[0]
 
+        self.sharding_device = sharding_device
         if self.sharding_device is not None:
             self.observed_pinn_in = jax.lax.with_sharding_constraint(
                 self.observed_pinn_in, self.sharding_device
@@ -126,7 +174,9 @@ class DataGeneratorObservations(AbstractDataGenerator):
         self.key, _ = jax.random.split(self.key, 2)  # to make it equivalent to
         # the call to _reset_batch_idx_and_permute in legacy DG
 
-    def _get_operands(self) -> tuple[Key, Int[Array, " n"], int, int | None, None]:
+    def _get_operands(
+        self,
+    ) -> tuple[PRNGKeyArray, Int[Array, " n"], int, int | None, None]:
         return (
             self.key,
             self.indices,
@@ -137,17 +187,19 @@ class DataGeneratorObservations(AbstractDataGenerator):
 
     def obs_batch(
         self,
-    ) -> tuple[DataGeneratorObservations, ObsBatchDict]:
+    ) -> tuple[Self, ObsBatchDict]:
         """
         Return an update DataGeneratorObservations instance and an ObsBatchDict
         """
         if self.obs_batch_size is None or self.obs_batch_size == self.n:
             # Avoid unnecessary reshuffling
-            return self, {
-                "pinn_in": self.observed_pinn_in,
-                "val": self.observed_values,
-                "eq_params": self.observed_eq_params,
-            }
+            return self, ObsBatchDict(
+                {
+                    "pinn_in": self.observed_pinn_in,
+                    "val": self.observed_values,
+                    "eq_params": self.observed_eq_params,
+                }
+            )
 
         new_attributes = _reset_or_increment(
             self.curr_idx + self.obs_batch_size,
@@ -157,7 +209,9 @@ class DataGeneratorObservations(AbstractDataGenerator):
             # handled above
         )
         new = eqx.tree_at(
-            lambda m: (m.key, m.indices, m.curr_idx), self, new_attributes
+            lambda m: (m.key, m.indices, m.curr_idx),  # type: ignore
+            self,
+            new_attributes,
         )
 
         minib_indices = jax.lax.dynamic_slice(
@@ -174,7 +228,7 @@ class DataGeneratorObservations(AbstractDataGenerator):
                 new.observed_values, minib_indices, unique_indices=True, axis=0
             ),
             "eq_params": jax.tree_util.tree_map(
-                lambda a: jnp.take(a, minib_indices, unique_indices=True, axis=0),
+                lambda a: jnp.take(a, minib_indices, unique_indices=True, axis=0),  # type: ignore
                 new.observed_eq_params,
             ),
         }
@@ -182,7 +236,7 @@ class DataGeneratorObservations(AbstractDataGenerator):
 
     def get_batch(
         self,
-    ) -> tuple[DataGeneratorObservations, ObsBatchDict]:
+    ) -> tuple[Self, ObsBatchDict]:
         """
         Generic method to return a batch
         """

jinns/data/_DataGeneratorParameter.py

@@ -5,12 +5,23 @@ Define the DataGenerators modules
 from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
+from typing import Self
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Key, Array, Float
+from jaxtyping import PRNGKeyArray, Array, Float
 from jinns.data._utils import _reset_or_increment
 from jinns.data._AbstractDataGenerator import AbstractDataGenerator
+from jinns.utils._DictToModuleMeta import DictToModuleMeta
+
+
+class DGParams(metaclass=DictToModuleMeta):
+    """
+    However, static type checkers cannot know that DGParams inherit from
+    eqx.Module and explicit casting to the latter class will be needed
+    """
+
+    pass
 
 
 class DataGeneratorParameter(AbstractDataGenerator):
@@ -21,9 +32,8 @@ class DataGeneratorParameter(AbstractDataGenerator):
 
     Parameters
     ----------
-    keys : Key | dict[str, Key]
-        Jax random key to sample new time points and to shuffle batches
-        or a dict of Jax random keys with key entries from param_ranges
+    key : PRNGKeyArray
+        Jax random key to sample new time points and to shuffle batches.
     n : int
         The number of total points that will be divided in
         batches. Batches are made so that each data point is seen only
@@ -58,71 +68,86 @@ class DataGeneratorParameter(AbstractDataGenerator):
         Defaults to None.
     """
 
-    keys: Key | dict[str, Key]
+    key: PRNGKeyArray
     n: int = eqx.field(static=True)
-    param_batch_size: int | None = eqx.field(static=True, default=None)
-    param_ranges: dict[str, tuple[Float, Float]] = eqx.field(
-        static=True, default_factory=lambda: {}
-    )
-    method: str = eqx.field(static=True, default="uniform")
-    user_data: dict[str, Float[Array, " n"]] | None = eqx.field(
-        default_factory=lambda: {}
-    )
+    param_batch_size: int | None = eqx.field(static=True)
+    param_ranges: dict[str, tuple[Float, Float]] = eqx.field(static=True)
+    method: str = eqx.field(static=True)
+    user_data: dict[str, Float[Array, " n"]]
+
+    # --- Below fields are not passed as arguments to __init__
+    _all_params_keys: set[str] = eqx.field(init=False, static=True)
+    curr_param_idx: eqx.Module | None = eqx.field(init=False)
+    param_n_samples: eqx.Module = eqx.field(init=False)
 
-    curr_param_idx: dict[str, int] = eqx.field(init=False)
-    param_n_samples: dict[str, Array] = eqx.field(init=False)
+    def __init__(
+        self,
+        *,
+        key: PRNGKeyArray,
+        n: int,
+        param_batch_size: int | None,
+        param_ranges: dict[str, tuple[Float, Float]] = {},
+        method: str = "uniform",
+        user_data: dict[str, Float[Array, " n"]] = {},
+    ):
+        self.key = key
+        self.n = n
+        self.param_batch_size = param_batch_size
+        self.param_ranges = param_ranges
+        self.method = method
+        self.user_data = user_data
+
+        _all_keys = set().union(self.param_ranges, self.user_data)
+        self._all_params_keys = _all_keys
 
-    def __post_init__(self):
-        if self.user_data is None:
-            self.user_data = {}
-        if self.param_ranges is None:
-            self.param_ranges = {}
         if self.param_batch_size is not None and self.n < self.param_batch_size:
             raise ValueError(
                 f"Number of data points ({self.n}) is smaller than the"
                 f"number of batch points ({self.param_batch_size})."
             )
-        if not isinstance(self.keys, dict):
-            all_keys = set().union(self.param_ranges, self.user_data)
-            self.keys = dict(zip(all_keys, jax.random.split(self.keys, len(all_keys))))
+
+        # NOTE from jinns > v1.5.1 we work with eqx.Module
+        # because eq_params is not a dict anymore.
+        # We have to use a different class from the publicly exposed EqParams
+        # because fields(EqParams) are not necessarily all present in the
+        # datagenerator, which would cause eqx.Module to error.
+
+        # 1) Call self.generate_data() to generate a dictionnary that merges the scattered data between `user_data` and `param_ranges`
+        self.key, _param_n_samples = self.generate_data(self.key)
+
+        # 2) Use the dictionnary to populate the field of the eqx.Module.
+        self.param_n_samples = DGParams(_param_n_samples, "DGParams")
 
         if self.param_batch_size is None:
-            self.curr_param_idx = None  # type: ignore
+            self.curr_param_idx = None
         else:
-            self.curr_param_idx = {}
-            for k in self.keys.keys():
-                self.curr_param_idx[k] = self.n + self.param_batch_size
-                # to be sure there is a shuffling at first get_batch()
+            curr_idx = self.n + self.param_batch_size
+            param_keys_and_curr_idx = {k: curr_idx for k in self._all_params_keys}
 
-        # The call to self.generate_data() creates
-        # the dict self.param_n_samples and then we will only use this one
-        # because it merges the scattered data between `user_data` and
-        # `param_ranges`
-        self.keys, self.param_n_samples = self.generate_data(self.keys)
+            self.curr_param_idx = DGParams(param_keys_and_curr_idx)
 
     def generate_data(
-        self, keys: dict[str, Key]
-    ) -> tuple[dict[str, Key], dict[str, Float[Array, " n"]]]:
+        self, key: PRNGKeyArray
+    ) -> tuple[PRNGKeyArray, dict[str, Float[Array, " n 1"]]]:
         """
         Generate parameter samples, either through generation
         or using user-provided data.
         """
         param_n_samples = {}
 
-        all_keys = set().union(
-            self.param_ranges,
-            self.user_data,  # type: ignore this has been handled in post_init
-        )
-        for k in all_keys:
+        # Some of the subkeys might not be used cause of user-provided data.
+        # This is not a big deal and simpler like that.
+        key, *subkeys = jax.random.split(key, len(self._all_params_keys) + 1)
+        for i, k in enumerate(self._all_params_keys):
             if self.user_data and k in self.user_data.keys():
-                if self.user_data[k].shape == (self.n, 1):
-                    param_n_samples[k] = self.user_data[k]
-                if self.user_data[k].shape == (self.n,):
-                    param_n_samples[k] = self.user_data[k][:, None]
-                else:
-                    raise ValueError(
+                try:
+                    param_n_samples[k] = self.user_data[k].reshape((self.n, 1))
+                except TypeError:
+                    shape = self.user_data[k].shape
+                    raise TypeError(
                         "Wrong shape for user provided parameters"
-                        f" in user_data dictionary at key='{k}'"
+                        f" in user_data dictionary at key='{k}' got {shape} "
+                        f"and expected {(self.n, 1)}."
                     )
             else:
                 if self.method == "grid":
@@ -132,75 +157,101 @@ class DataGeneratorParameter(AbstractDataGenerator):
                     param_n_samples[k] = jnp.arange(xmin, xmax, partial)[:, None]
                 elif self.method == "uniform":
                     xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
-                    keys[k], subkey = jax.random.split(keys[k], 2)
                     param_n_samples[k] = jax.random.uniform(
-                        subkey, shape=(self.n, 1), minval=xmin, maxval=xmax
+                        subkeys[i], shape=(self.n, 1), minval=xmin, maxval=xmax
                     )
                 else:
                     raise ValueError("Method " + self.method + " is not implemented.")
 
-        return keys, param_n_samples
-
-    def _get_param_operands(
-        self, k: str
-    ) -> tuple[Key, Float[Array, " n"], int, int | None, None]:
-        return (
-            self.keys[k],
-            self.param_n_samples[k],
-            self.curr_param_idx[k],
-            self.param_batch_size,
-            None,
-        )
+        return key, param_n_samples
 
-    def param_batch(self):
+    def param_batch(self) -> tuple[Self, eqx.Module]:
         """
-        Return a dictionary with batches of parameters
-        If all the batches have been seen, we reshuffle them,
-        otherwise we just return the next unseen batch.
+        Return an `eqx.Module` with batches of parameters at its leafs.
+        If all the batches have been seen, we reshuffle them (or rather
+        their indices), otherwise we just return the next unseen batch.
         """
 
         if self.param_batch_size is None or self.param_batch_size == self.n:
+            # Full batch mode: nothing to do.
             return self, self.param_n_samples
+        else:
+
+            def _reset_or_increment_wrapper(
+                param_k: Array, idx_k: int, key_k: PRNGKeyArray
+            ):
+                everything_but_key = _reset_or_increment(
+                    idx_k + self.param_batch_size,  # type: ignore
+                    self.n,
+                    (key_k, param_k, idx_k, self.param_batch_size, None),  # type: ignore
+                )[1:]
+                return everything_but_key
 
-        def _reset_or_increment_wrapper(param_k, idx_k, key_k):
-            return _reset_or_increment(
-                idx_k + self.param_batch_size,
-                self.n,
-                (key_k, param_k, idx_k, self.param_batch_size, None),  # type: ignore
-                # ignore since the case self.param_batch_size is None has been
-                # handled above
+            new_key, *subkeys = jax.random.split(
+                self.key, len(self._all_params_keys) + 1
             )
 
-        res = jax.tree_util.tree_map(
-            _reset_or_increment_wrapper,
-            self.param_n_samples,
-            self.curr_param_idx,
-            self.keys,
-        )
-        # we must transpose the pytrees because keys are merged in res
-        # https://jax.readthedocs.io/en/latest/jax-101/05.1-pytrees.html#transposing-trees
-        new_attributes = jax.tree_util.tree_transpose(
-            jax.tree_util.tree_structure(self.keys),
-            jax.tree_util.tree_structure([0, 0, 0]),
-            res,
-        )
-
-        new = eqx.tree_at(
-            lambda m: (m.keys, m.param_n_samples, m.curr_param_idx),
-            self,
-            new_attributes,
-        )
-
-        return new, jax.tree_util.tree_map(
-            lambda p, q: jax.lax.dynamic_slice(
-                p, start_indices=(q, 0), slice_sizes=(new.param_batch_size, 1)
-            ),
-            new.param_n_samples,
-            new.curr_param_idx,
-        )
-
-    def get_batch(self):
+            # From PRNGKeyArray to a pytree of keys with adequate structure
+            subkeys = jax.tree.unflatten(
+                jax.tree.structure(self.param_n_samples), subkeys
+            )
+
+            res = jax.tree.map(
+                _reset_or_increment_wrapper,
+                self.param_n_samples,
+                self.curr_param_idx,
+                subkeys,
+            )
+            # we must transpose the pytrees because both params and curr_idx # are merged in res
+            # https://jax.readthedocs.io/en/latest/jax-101/05.1-pytrees.html#transposing-trees
+
+            new_attributes = jax.tree.transpose(
+                jax.tree.structure(self.param_n_samples),
+                jax.tree.structure([0, 0]),
+                res,
+            )
+
+            new = eqx.tree_at(
+                lambda m: (m.param_n_samples, m.curr_param_idx),
+                self,
+                new_attributes,
+            )
+
+            new = eqx.tree_at(lambda m: m.key, new, new_key)
+
+            return new, jax.tree_util.tree_map(
+                lambda p, q: jax.lax.dynamic_slice(
+                    p, start_indices=(q, 0), slice_sizes=(new.param_batch_size, 1)
+                ),
+                new.param_n_samples,
+                new.curr_param_idx,
+            )
+
+    def get_batch(self) -> tuple[Self, eqx.Module]:
         """
         Generic method to return a batch
         """
         return self.param_batch()
+
+
+if __name__ == "__main__":
+    key = jax.random.PRNGKey(2)
+    key, subkey = jax.random.split(key)
+
+    n = 64
+    param_batch_size = 32
+    method = "uniform"
+    param_ranges = {"theta": (10.0, 11.0)}
+    user_data = {"nu": jnp.ones((n, 1))}
+
+    x = DataGeneratorParameter(
+        key=subkey,
+        n=n,
+        param_batch_size=param_batch_size,
+        param_ranges=param_ranges,
+        method=method,
+        user_data=user_data,
+    )
+    print(x.key)
+    x, batch = x.get_batch()
+    print(x.key)