jinns 0.9.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

jinns/data/_DataGenerators.py

@@ -1,52 +1,49 @@
+# pylint: disable=unsubscriptable-object
 """
-DataGenerators to generate batches of points in space, time and more
+Define the DataGeneratorODE equinox module
 """
-
-from typing import NamedTuple
-from jax.typing import ArrayLike
+from __future__ import (
+    annotations,
+)  # https://docs.python.org/3/library/typing.html#constant
+
+from typing import TYPE_CHECKING, Dict
+from dataclasses import InitVar
+import equinox as eqx
+import jax
 import jax.numpy as jnp
-from jax import random
-from jax.tree_util import register_pytree_node_class
-import jax.lax
-
-
-class ODEBatch(NamedTuple):
-    temporal_batch: ArrayLike
-    param_batch_dict: dict = None
-    obs_batch_dict: dict = None
-
-
-class PDENonStatioBatch(NamedTuple):
-    times_x_inside_batch: ArrayLike
-    times_x_border_batch: ArrayLike
-    param_batch_dict: dict = None
-    obs_batch_dict: dict = None
-
+from jaxtyping import Key, Int, PyTree, Array, Float, Bool
+from jinns.data._Batchs import *
 
-class PDEStatioBatch(NamedTuple):
-    inside_batch: ArrayLike
-    border_batch: ArrayLike
-    param_batch_dict: dict = None
-    obs_batch_dict: dict = None
+if TYPE_CHECKING:
+    from jinns.utils._types import *
 
 
-def append_param_batch(batch, param_batch_dict):
+def append_param_batch(batch: AnyBatch, param_batch_dict: dict) -> AnyBatch:
     """
     Utility function that fill the param_batch_dict of a batch object with a
     param_batch_dict
     """
-    return batch._replace(param_batch_dict=param_batch_dict)
+    return eqx.tree_at(
+        lambda m: m.param_batch_dict,
+        batch,
+        param_batch_dict,
+        is_leaf=lambda x: x is None,
+    )
 
 
-def append_obs_batch(batch, obs_batch_dict):
+def append_obs_batch(batch: AnyBatch, obs_batch_dict: dict) -> AnyBatch:
     """
     Utility function that fill the obs_batch_dict of a batch object with a
     obs_batch_dict
     """
-    return batch._replace(obs_batch_dict=obs_batch_dict)
+    return eqx.tree_at(
+        lambda m: m.obs_batch_dict, batch, obs_batch_dict, is_leaf=lambda x: x is None
+    )
 
 
-def make_cartesian_product(b1, b2):
+def make_cartesian_product(
+    b1: Float[Array, "batch_size dim1"], b2: Float[Array, "batch_size dim2"]
+) -> Float[Array, "(batch_size*batch_size) (dim1+dim2)"]:
     """
     Create the cartesian product of a time and a border omega batches
     by tiling and repeating
@@ -58,29 +55,39 @@ def make_cartesian_product(b1, b2):
     return jnp.concatenate([b1, b2], axis=1)
 
 
-def _reset_batch_idx_and_permute(operands):
+def _reset_batch_idx_and_permute(
+    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]]
+) -> tuple[Key, Float[Array, "n dimension"], Int]:
     key, domain, curr_idx, _, p = operands
     # resetting counter
     curr_idx = 0
     # reshuffling
-    key, subkey = random.split(key)
+    key, subkey = jax.random.split(key)
     # domain = random.permutation(subkey, domain, axis=0, independent=False)
     # we want that permutation = choice when p=None
     # otherwise p is used to avoid collocation points not in nt_start
-    domain = random.choice(subkey, domain, shape=(domain.shape[0],), replace=False, p=p)
+    domain = jax.random.choice(
+        subkey, domain, shape=(domain.shape[0],), replace=False, p=p
+    )
 
     # return updated
     return (key, domain, curr_idx)
 
 
-def _increment_batch_idx(operands):
+def _increment_batch_idx(
+    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]]
+) -> tuple[Key, Float[Array, "n dimension"], Int]:
     key, domain, curr_idx, batch_size, _ = operands
     # simply increases counter and get the batch
     curr_idx += batch_size
     return (key, domain, curr_idx)
 
 
-def _reset_or_increment(bend, n_eff, operands):
+def _reset_or_increment(
+    bend: Int,
+    n_eff: Int,
+    operands: tuple[Key, Float[Array, "n dimension"], Int, None, Float[Array, "n"]],
+) -> tuple[Key, Float[Array, "n dimension"], Int]:
     """
     Factorize the code of the jax.lax.cond which checks if we have seen all the
     batches in an epoch
@@ -109,15 +116,18 @@ def _reset_or_increment(bend, n_eff, operands):
     )
 
 
-def _check_and_set_rar_parameters(rar_parameters, n, n_start):
+def _check_and_set_rar_parameters(
+    rar_parameters: dict, n: Int, n_start: Int
+) -> tuple[Int, Float[Array, "n"], Int, Int]:
     if rar_parameters is not None and n_start is None:
         raise ValueError(
-            f"n_start or/and nt_start must be provided in the context of RAR sampling scheme, {n_start} was provided"
+            "nt_start must be provided in the context of RAR sampling scheme"
         )
+
     if rar_parameters is not None:
         # Default p is None. However, in the RAR sampling scheme we use 0
         # probability to specify non-used collocation points (i.e. points
-        # above n_start). Thus, p is a vector of probability of shape (n, 1).
+        # above nt_start). Thus, p is a vector of probability of shape (nt, 1).
         p = jnp.zeros((n,))
         p = p.at[:n_start].set(1 / n_start)
         # set internal counter for the number of gradient steps since the
@@ -129,6 +139,7 @@ def _check_and_set_rar_parameters(rar_parameters, n, n_start):
         # have been added
         rar_iter_nb = 0
     else:
+        n_start = n
         p = None
         rar_iter_from_last_sampling = None
         rar_iter_nb = None
@@ -136,109 +147,102 @@ def _check_and_set_rar_parameters(rar_parameters, n, n_start):
     return n_start, p, rar_iter_from_last_sampling, rar_iter_nb
 
 
-#####################################################
-# DataGenerator for ODE : only returns time_batches
-#####################################################
-
-
-@register_pytree_node_class
-class DataGeneratorODE:
+class DataGeneratorODE(eqx.Module):
     """
     A class implementing data generator object for ordinary differential equations.
 
-
-    **Note:** DataGeneratorODE is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+    key : Key
+        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
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    tmin : float
+        The minimum value of the time domain to consider
+    tmax : float
+        The maximum value of the time domain to consider
+    temporal_batch_size : int
+        The size of the batch of randomly selected points among
+        the `nt` points.
+    method : str, default="uniform"
+        Either `grid` or `uniform`, default is `uniform`.
+        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 : Dict[str, Int], default=None
+        Default to None: do not use Residual Adaptative Resampling.
+        Otherwise a dictionary with keys. `start_iter`: the iteration at
+        which we start the RAR sampling scheme (we first have a burn in
+        period). `update_rate`: the number of gradient steps taken between
+        each appending of collocation points in the RAR algo.
+        `sample_size`: the size of the sample from which we will select new
+        collocation points. `selected_sample_size_times`: the number of selected
+        points from the sample to be added to the current collocation
+        points
+    nt_start : 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
+        nt_start = nt and this is hidden from the user.
+        In RAR, nt_start
+        then corresponds to the initial number of points we train the PINN.
     """
 
-    def __init__(
-        self,
-        key,
-        nt,
-        tmin,
-        tmax,
-        temporal_batch_size,
-        method="uniform",
-        rar_parameters=None,
-        nt_start=None,
-        data_exists=False,
-    ):
-        """
-        Parameters
-        ----------
-        key
-            Jax random key to sample new time points and to shuffle batches
-        nt
-            An integer. The number of total time points that will be divided in
-            batches. Batches are made so that each data point is seen only
-            once during 1 epoch.
-        tmin
-            A float. The minimum value of the time domain to consider
-        tmax
-            A float. The maximum value of the time domain to consider
-        temporal_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `nt` points.
-        method
-            Either `grid` or `uniform`, default is `uniform`.
-            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
-            Default to None: do not use Residual Adaptative Resampling.
-            Otherwise a dictionary with keys. `start_iter`: the iteration at
-            which we start the RAR sampling scheme (we first have a burn in
-            period). `update_every`: the number of gradient steps taken between
-            each appending of collocation points in the RAR algo.
-            `sample_size_times`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size_times`: the number of selected
-            points from the sample to be added to the current collocation
-            points
-            "DeepXDE: A deep learning library for solving differential
-            equations", L. Lu, SIAM Review, 2021
-        nt_start
-            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
-            nt_start = nt and this is hidden from the user.
-            In RAR, nt_start
-            then corresponds to the initial number of points we train the PINN.
-        data_exists
-            Must be left to `False` when created by the user. Avoids the
-            regeneration of the `nt` time points at each pytree flattening and
-            unflattening.
-        """
-        self.data_exists = data_exists
-        self._key = key
-        self.nt = nt
-        self.tmin = tmin
-        self.tmax = tmax
-        self.temporal_batch_size = temporal_batch_size
-        self.method = method
-        self.rar_parameters = rar_parameters
-
-        # Set-up for RAR (if used)
+    key: Key
+    nt: Int
+    tmin: Float
+    tmax: Float
+    temporal_batch_size: Int = eqx.field(static=True)  # static cause used as a
+    # shape in jax.lax.dynamic_slice
+    method: str = eqx.field(static=True, default_factory=lambda: "uniform")
+    rar_parameters: Dict[str, Int] = None
+    nt_start: Int = eqx.field(static=True, default=None)
+
+    # all the init=False fields are set in __post_init__, even after a _replace
+    # or eqx.tree_at __post_init__ is called
+    p_times: Float[Array, "nt"] = eqx.field(init=False)
+    rar_iter_from_last_sampling: Int = eqx.field(init=False)
+    rar_iter_nb: Int = eqx.field(init=False)
+    curr_time_idx: Int = eqx.field(init=False)
+    times: Float[Array, "nt"] = eqx.field(init=False)
+
+    def __post_init__(self):
         (
             self.nt_start,
             self.p_times,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
-        ) = _check_and_set_rar_parameters(rar_parameters, n=nt, n_start=nt_start)
-
-        if not self.data_exists:
-            # Useful when using a lax.scan with pytree
-            # Optionally can tell JAX not to re-generate data
-            self.curr_time_idx = 0
-            self.generate_time_data()
-            self._key, self.times, _ = _reset_batch_idx_and_permute(
-                self._get_time_operands()
-            )
-
-    def sample_in_time_domain(self, n_samples):
-        self._key, subkey = random.split(self._key, 2)
-        return random.uniform(subkey, (n_samples,), minval=self.tmin, maxval=self.tmax)
+        ) = _check_and_set_rar_parameters(self.rar_parameters, self.nt, self.nt_start)
+
+        self.curr_time_idx = jnp.iinfo(jnp.int32).max - self.temporal_batch_size - 1
+        # to be sure there is a
+        # shuffling at first get_batch() we do not call
+        # _reset_batch_idx_and_permute in __init__ or __post_init__ because it
+        # would return a copy of self and we have not investigate what would
+        # happen
+        # NOTE the (- self.temporal_batch_size - 1) because otherwise when computing
+        # `bend` we overflow the max int32 with unwanted behaviour
+
+        self.key, self.times = self.generate_time_data(self.key)
+        # Note that, here, in __init__ (and __post_init__), this is the
+        # only place where self assignment are authorized so we do the
+        # above way for the key. Note that one of the motivation to return the
+        # key from generate_*_data is to easily align key with legacy
+        # DataGenerators to use same unit tests
+
+    def sample_in_time_domain(
+        self, key: Key, sample_size: Int = None
+    ) -> Float[Array, "nt"]:
+        return jax.random.uniform(
+            key,
+            (self.nt if sample_size is None else sample_size,),
+            minval=self.tmin,
+            maxval=self.tmax,
+        )
 
-    def generate_time_data(self):
+    def generate_time_data(self, key: Key) -> tuple[Key, Float[Array, "nt"]]:
         """
         Construct a complete set of `self.nt` time points according to the
         specified `self.method`
@@ -246,24 +250,28 @@ class DataGeneratorODE:
         Note that self.times has always size self.nt and not self.nt_start, even
         in RAR scheme, we must allocate all the collocation points
         """
+        key, subkey = jax.random.split(self.key)
         if self.method == "grid":
-            self.partial_times = (self.tmax - self.tmin) / self.nt
-            self.times = jnp.arange(self.tmin, self.tmax, self.partial_times)
-        elif self.method == "uniform":
-            self.times = self.sample_in_time_domain(self.nt)
-        else:
-            raise ValueError("Method " + self.method + " is not implemented.")
+            partial_times = (self.tmax - self.tmin) / self.nt
+            return key, jnp.arange(self.tmin, self.tmax, partial_times)
+        if self.method == "uniform":
+            return key, self.sample_in_time_domain(subkey)
+        raise ValueError("Method " + self.method + " is not implemented.")
 
-    def _get_time_operands(self):
+    def _get_time_operands(
+        self,
+    ) -> tuple[Key, Float[Array, "nt"], Int, Int, Float[Array, "nt"]]:
         return (
-            self._key,
+            self.key,
             self.times,
             self.curr_time_idx,
             self.temporal_batch_size,
             self.p_times,
         )
 
-    def temporal_batch(self):
+    def temporal_batch(
+        self,
+    ) -> tuple["DataGeneratorODE", Float[Array, "temporal_batch_size"]]:
         """
         Return a batch of time points. If all the batches have been seen, we
         reshuffle them, otherwise we just return the next unseen batch.
@@ -275,210 +283,142 @@ class DataGeneratorODE:
         if self.rar_parameters is not None:
             nt_eff = (
                 self.nt_start
-                + self.rar_iter_nb * self.rar_parameters["selected_sample_size_omega"]
+                + self.rar_iter_nb * self.rar_parameters["selected_sample_size_times"]
             )
         else:
             nt_eff = self.nt
-        (self._key, self.times, self.curr_time_idx) = _reset_or_increment(
-            bend, nt_eff, self._get_time_operands()
+
+        new_attributes = _reset_or_increment(bend, nt_eff, self._get_time_operands())
+        new = eqx.tree_at(
+            lambda m: (m.key, m.times, m.curr_time_idx), self, new_attributes
         )
 
         # commands below are equivalent to
         # return self.times[i:(i+t_batch_size)]
         # start indices can be dynamic be the slice shape is fixed
-        return jax.lax.dynamic_slice(
-            self.times,
-            start_indices=(self.curr_time_idx,),
-            slice_sizes=(self.temporal_batch_size,),
+        return new, jax.lax.dynamic_slice(
+            new.times,
+            start_indices=(new.curr_time_idx,),
+            slice_sizes=(new.temporal_batch_size,),
         )
 
-    def get_batch(self):
+    def get_batch(self) -> tuple["DataGeneratorODE", ODEBatch]:
         """
         Generic method to return a batch. Here we call `self.temporal_batch()`
         """
-        return ODEBatch(temporal_batch=self.temporal_batch())
-
-    def tree_flatten(self):
-        children = (
-            self._key,
-            self.times,
-            self.curr_time_idx,
-            self.tmin,
-            self.tmax,
-            self.p_times,
-            self.rar_iter_from_last_sampling,
-            self.rar_iter_nb,
-        )  # arrays / dynamic values
-        aux_data = {
-            k: vars(self)[k]
-            for k in [
-                "temporal_batch_size",
-                "method",
-                "nt",
-                "rar_parameters",
-                "nt_start",
-            ]
-        }  # static values
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        """
-        **Note:** When reconstructing the class, we force ``data_exists=True``
-        in order not to re-generate the data at each flattening and
-        unflattening that happens e.g. during the gradient descent in the
-        optimization process
-        """
-        (
-            key,
-            times,
-            curr_time_idx,
-            tmin,
-            tmax,
-            p_times,
-            rar_iter_from_last_sampling,
-            rar_iter_nb,
-        ) = children
-        obj = cls(
-            key=key,
-            data_exists=True,
-            tmin=tmin,
-            tmax=tmax,
-            **aux_data,
-        )
-        obj.times = times
-        obj.curr_time_idx = curr_time_idx
-        obj.p_times = p_times
-        obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
-        obj.rar_iter_nb = rar_iter_nb
-        return obj
-
-
-##########################################
-# Data Generator for PDE in stationnary
-# and non-stationnary cases
-##########################################
-
-
-class DataGeneratorPDEAbstract:
-    """generic skeleton class for a PDE data generator"""
+        new, temporal_batch = self.temporal_batch()
+        return new, ODEBatch(temporal_batch=temporal_batch)
 
-    def __init__(self, data_exists=False) -> None:
-        # /!\ WARNING /!\: an-end user should never create an object
-        # with data_exists=True. Or else generate_data() won't be called.
-        # Useful when using a lax.scan with a DataGenerator in the carry
-        # It tells JAX not to re-generate data in the __init__()
-        self.data_exists = data_exists
 
-
-@register_pytree_node_class
-class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
-    """
+class CubicMeshPDEStatio(eqx.Module):
+    r"""
     A class implementing data generator object for stationary partial
     differential equations.
 
-
-    **Note:** CubicMeshPDEStatio is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+    key : Key
+        Jax random key to sample new time points and to shuffle batches
+    n : Int
+        The number of total $\Omega$ points that will be divided in
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    nb : Int | None
+        The total number of points in $\partial\Omega$.
+        Can be `None` not to lose performance generating the border
+        batch if they are not used
+    omega_batch_size : Int
+        The size of the batch of randomly selected points among
+        the `n` points.
+    omega_border_batch_size : Int | None
+        The size of the batch of points randomly selected
+        among the `nb` points.
+        Can be `None` not to lose performance generating the border
+        batch if they are not used
+    dim : Int
+        Dimension of $\Omega$ domain
+    min_pts : tuple[tuple[Float, Float], ...]
+        A tuple of minimum values of the domain along each dimension. For a sampling
+        in `n` dimension, this represents $(x_{1, min}, x_{2,min}, ...,
+        x_{n, min})$
+    max_pts : tuple[tuple[Float, Float], ...]
+        A tuple of maximum values of the domain along each dimension. For a sampling
+        in `n` dimension, this represents $(x_{1, max}, x_{2,max}, ...,
+        x_{n,max})$
+    method : str, default="uniform"
+        Either `grid` or `uniform`, default is `uniform`.
+        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 : Dict[str, Int], default=None
+        Default to None: do not use Residual Adaptative Resampling.
+        Otherwise a dictionary with keys. `start_iter`: the iteration at
+        which we start the RAR sampling scheme (we first have a burn in
+        period). `update_every`: the number of gradient steps taken between
+        each appending of collocation points in the RAR algo.
+        `sample_size_omega`: the size of the sample from which we will select new
+        collocation points. `selected_sample_size_omega`: the number of selected
+        points from the sample to be added to the current collocation
+        points
+    n_start : Int, default=None
+        Defaults to None. The effective size of n used at start time.
+        This value must be
+        provided when rar_parameters is not None. Otherwise we set internally
+        n_start = n and this is hidden from the user.
+        In RAR, n_start
+        then corresponds to the initial number of points we train the PINN.
     """
 
-    def __init__(
-        self,
-        key,
-        n,
-        nb,
-        omega_batch_size,
-        omega_border_batch_size,
-        dim,
-        min_pts,
-        max_pts,
-        method="grid",
-        rar_parameters=None,
-        n_start=None,
-        data_exists=False,
-    ):
-        r"""
-        Parameters
-        ----------
-        key
-            Jax random key to sample new time points and to shuffle batches
-        n
-            An integer. The number of total :math:`\Omega` points that will be divided in
-            batches. Batches are made so that each data point is seen only
-            once during 1 epoch.
-        nb
-            An integer. The total number of points in :math:`\partial\Omega`.
-            Can be `None` not to lose performance generating the border
-            batch if they are not used
-        omega_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `n` points.
-        omega_border_batch_size
-            An integer. The size of the batch of points randomly selected
-            among the `nb` points.
-            Can be `None` not to lose performance generating the border
-            batch if they are not used
-        dim
-            An integer. dimension of :math:`\Omega` domain
-        min_pts
-            A tuple of minimum values of the domain along each dimension. For a sampling
-            in `n` dimension, this represents :math:`(x_{1, min}, x_{2,min}, ...,
-            x_{n, min})`
-        max_pts
-            A tuple of maximum values of the domain along each dimension. For a sampling
-            in `n` dimension, this represents :math:`(x_{1, max}, x_{2,max}, ...,
-            x_{n,max})`
-        method
-            Either `grid` or `uniform`, default is `grid`.
-            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
-            Default to None: do not use Residual Adaptative Resampling.
-            Otherwise a dictionary with keys. `start_iter`: the iteration at
-            which we start the RAR sampling scheme (we first have a burn in
-            period). `update_every`: the number of gradient steps taken between
-            each appending of collocation points in the RAR algo.
-            `sample_size_omega`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size_omega`: the number of selected
-            points from the sample to be added to the current collocation
-            points
-            "DeepXDE: A deep learning library for solving differential
-            equations", L. Lu, SIAM Review, 2021
-        n_start
-            Defaults to None. The effective size of n used at start time.
-            This value must be
-            provided when rar_parameters is not None. Otherwise we set internally
-            n_start = n and this is hidden from the user.
-            In RAR, n_start
-            then corresponds to the initial number of points we train the PINN.
-        data_exists
-            Must be left to `False` when created by the user. Avoids the
-            regeneration of :math:`\Omega`, :math:`\partial\Omega` and
-            time points at each pytree flattening and unflattening.
-        """
-        super().__init__(data_exists=data_exists)
-        self.method = method
-        self._key = key
-        self.dim = dim
-        self.min_pts = min_pts
-        self.max_pts = max_pts
-        assert dim == len(min_pts) and isinstance(min_pts, tuple)
-        assert dim == len(max_pts) and isinstance(max_pts, tuple)
-        self.n = n
-        self.rar_parameters = rar_parameters
+    # kw_only in base class is motivated here: https://stackoverflow.com/a/69822584
+    key: Key = eqx.field(kw_only=True)
+    n: Int = eqx.field(kw_only=True)
+    nb: Int | None = eqx.field(kw_only=True)
+    omega_batch_size: Int = eqx.field(
+        kw_only=True, static=True
+    )  # static cause used as a
+    # shape in jax.lax.dynamic_slice
+    omega_border_batch_size: Int | None = eqx.field(
+        kw_only=True, static=True
+    )  # static cause used as a
+    # shape in jax.lax.dynamic_slice
+    dim: Int = eqx.field(kw_only=True, static=True)  # static cause used as a
+    # shape in jax.lax.dynamic_slice
+    min_pts: tuple[tuple[Float, Float], ...] = eqx.field(kw_only=True)
+    max_pts: tuple[tuple[Float, Float], ...] = eqx.field(kw_only=True)
+    method: str = eqx.field(
+        kw_only=True, static=True, default_factory=lambda: "uniform"
+    )
+    rar_parameters: Dict[str, Int] = eqx.field(kw_only=True, default=None)
+    n_start: Int = eqx.field(kw_only=True, default=None, static=True)
+
+    # all the init=False fields are set in __post_init__, even after a _replace
+    # or eqx.tree_at __post_init__ is called
+    p_omega: Float[Array, "n"] = eqx.field(init=False)
+    p_border: None = eqx.field(init=False)
+    rar_iter_from_last_sampling: Int = eqx.field(init=False)
+    rar_iter_nb: Int = eqx.field(init=False)
+    curr_omega_idx: Int = eqx.field(init=False)
+    curr_omega_border_idx: Int = eqx.field(init=False)
+    omega: Float[Array, "n dim"] = eqx.field(init=False)
+    omega_border: Float[Array, "1 2"] | Float[Array, "(nb//4) 2 4"] | None = eqx.field(
+        init=False
+    )
+
+    def __post_init__(self):
+        assert self.dim == len(self.min_pts) and isinstance(self.min_pts, tuple)
+        assert self.dim == len(self.max_pts) and isinstance(self.max_pts, tuple)
+
         (
             self.n_start,
             self.p_omega,
             self.rar_iter_from_last_sampling,
             self.rar_iter_nb,
-        ) = _check_and_set_rar_parameters(rar_parameters, n=n, n_start=n_start)
+        ) = _check_and_set_rar_parameters(self.rar_parameters, self.n, self.n_start)
 
         self.p_border = None  # no RAR sampling for border for now
 
-        self.omega_batch_size = omega_batch_size
-
-        if omega_border_batch_size is None:
+        # Special handling for the border batch
+        if self.omega_border_batch_size is None:
             self.nb = None
             self.omega_border_batch_size = None
         elif self.dim == 1:
@@ -491,48 +431,46 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             # ignored since borders of Omega are singletons.
             #  self.border_batch() will return [xmin, xmax]
         else:
-            if nb % (2 * self.dim) != 0 or nb < 2 * self.dim:
+            if self.nb % (2 * self.dim) != 0 or self.nb < 2 * self.dim:
                 raise ValueError(
                     "number of border point must be"
                     " a multiple of 2xd (the # of faces of a d-dimensional cube)"
                 )
-            if nb // (2 * self.dim) < omega_border_batch_size:
+            if self.nb // (2 * self.dim) < self.omega_border_batch_size:
                 raise ValueError(
                     "number of points per facets (nb//2*self.dim)"
                     " cannot be lower than border batch size"
                 )
-            self.nb = int((2 * self.dim) * (nb // (2 * self.dim)))
-            self.omega_border_batch_size = omega_border_batch_size
-
-        if not self.data_exists:
-            # Useful when using a lax.scan with pytree
-            # Optionally tells JAX not to re-generate data when re-building the
-            # object
-            self.curr_omega_idx = 0
-            self.curr_omega_border_idx = 0
-            self.generate_data()
-            self._key, self.omega, _ = _reset_batch_idx_and_permute(
-                self._get_omega_operands()
-            )
-            if self.omega_border is not None and self.dim > 1:
-                self._key, self.omega_border, _ = _reset_batch_idx_and_permute(
-                    self._get_omega_border_operands()
-                )
+            self.nb = int((2 * self.dim) * (self.nb // (2 * self.dim)))
 
-    def sample_in_omega_domain(self, n_samples):
+        self.curr_omega_idx = jnp.iinfo(jnp.int32).max - self.omega_batch_size - 1
+        # see explaination in DataGeneratorODE
+        if self.omega_border_batch_size is None:
+            self.curr_omega_border_idx = None
+        else:
+            self.curr_omega_border_idx = (
+                jnp.iinfo(jnp.int32).max - self.omega_border_batch_size - 1
+            )
+        # key, subkey = jax.random.split(self.key)
+        # self.key = key
+        self.key, self.omega, self.omega_border = self.generate_data(self.key)
+        # see explaination in DataGeneratorODE for the key
+
+    def sample_in_omega_domain(
+        self, keys: Key, sample_size: Int = None
+    ) -> Float[Array, "n dim"]:
+        sample_size = self.n if sample_size is None else sample_size
         if self.dim == 1:
             xmin, xmax = self.min_pts[0], self.max_pts[0]
-            self._key, subkey = random.split(self._key, 2)
-            return random.uniform(
-                subkey, shape=(n_samples, 1), minval=xmin, maxval=xmax
+            return jax.random.uniform(
+                keys, shape=(sample_size, 1), minval=xmin, maxval=xmax
             )
-        keys = random.split(self._key, self.dim + 1)
-        self._key = keys[0]
+        # keys = jax.random.split(key, self.dim)
         return jnp.concatenate(
             [
-                random.uniform(
-                    keys[i + 1],
-                    (n_samples, 1),
+                jax.random.uniform(
+                    keys[i],
+                    (sample_size, 1),
                     minval=self.min_pts[i],
                     maxval=self.max_pts[i],
                 )
@@ -541,7 +479,9 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             axis=-1,
         )
 
-    def sample_in_omega_border_domain(self, n_samples):
+    def sample_in_omega_border_domain(
+        self, keys: Key
+    ) -> Float[Array, "1 2"] | Float[Array, "(nb//4) 2 4"] | None:
         if self.omega_border_batch_size is None:
             return None
         if self.dim == 1:
@@ -553,15 +493,12 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             # TODO : find a general & efficient way to sample from the border
             # (facets) of the hypercube in general dim.
 
-            facet_n = n_samples // (2 * self.dim)
-            keys = random.split(self._key, 5)
-            self._key = keys[0]
-            subkeys = keys[1:]
+            facet_n = self.nb // (2 * self.dim)
             xmin = jnp.hstack(
                 [
                     self.min_pts[0] * jnp.ones((facet_n, 1)),
-                    random.uniform(
-                        subkeys[0],
+                    jax.random.uniform(
+                        keys[0],
                         (facet_n, 1),
                         minval=self.min_pts[1],
                         maxval=self.max_pts[1],
@@ -571,8 +508,8 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             xmax = jnp.hstack(
                 [
                     self.max_pts[0] * jnp.ones((facet_n, 1)),
-                    random.uniform(
-                        subkeys[1],
+                    jax.random.uniform(
+                        keys[1],
                         (facet_n, 1),
                         minval=self.min_pts[1],
                         maxval=self.max_pts[1],
@@ -581,8 +518,8 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             )
             ymin = jnp.hstack(
                 [
-                    random.uniform(
-                        subkeys[2],
+                    jax.random.uniform(
+                        keys[2],
                         (facet_n, 1),
                         minval=self.min_pts[0],
                         maxval=self.max_pts[0],
@@ -592,8 +529,8 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             )
             ymax = jnp.hstack(
                 [
-                    random.uniform(
-                        subkeys[3],
+                    jax.random.uniform(
+                        keys[3],
                         (facet_n, 1),
                         minval=self.min_pts[0],
                         maxval=self.max_pts[0],
@@ -607,54 +544,71 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
             + f"implemented yet. You are asking for generation in dimension d={self.dim}."
         )
 
-    def generate_data(self):
+    def generate_data(self, key: Key) -> tuple[
+        Key,
+        Float[Array, "n dim"],
+        Float[Array, "1 2"] | Float[Array, "(nb//4) 2 4"] | None,
+    ]:
         r"""
-        Construct a complete set of `self.n` :math:`\Omega` points according to the
+        Construct a complete set of `self.n` $\Omega$ points according to the
         specified `self.method`. Also constructs a complete set of `self.nb`
-        :math:`\partial\Omega` points if `self.omega_border_batch_size` is not
+        $\partial\Omega$ points if `self.omega_border_batch_size` is not
         `None`. If the latter is `None` we set `self.omega_border` to `None`.
         """
-
         # Generate Omega
         if self.method == "grid":
             if self.dim == 1:
                 xmin, xmax = self.min_pts[0], self.max_pts[0]
-                self.partial = (xmax - xmin) / self.n
+                partial = (xmax - xmin) / self.n
                 # shape (n, 1)
-                self.omega = jnp.arange(xmin, xmax, self.partial)[:, None]
+                omega = jnp.arange(xmin, xmax, partial)[:, None]
             else:
-                self.partials = [
+                partials = [
                     (self.max_pts[i] - self.min_pts[i]) / jnp.sqrt(self.n)
                     for i in range(self.dim)
                 ]
                 xyz_ = jnp.meshgrid(
                     *[
-                        jnp.arange(self.min_pts[i], self.max_pts[i], self.partials[i])
+                        jnp.arange(self.min_pts[i], self.max_pts[i], partials[i])
                         for i in range(self.dim)
                     ]
                 )
                 xyz_ = [a.reshape((self.n, 1)) for a in xyz_]
-                self.omega = jnp.concatenate(xyz_, axis=-1)
+                omega = jnp.concatenate(xyz_, axis=-1)
         elif self.method == "uniform":
-            self.omega = self.sample_in_omega_domain(self.n)
+            if self.dim == 1:
+                key, subkeys = jax.random.split(key, 2)
+            else:
+                key, *subkeys = jax.random.split(key, self.dim + 1)
+            omega = self.sample_in_omega_domain(subkeys)
         else:
             raise ValueError("Method " + self.method + " is not implemented.")
 
         # Generate border of omega
-        self.omega_border = self.sample_in_omega_border_domain(self.nb)
+        if self.dim == 2 and self.omega_border_batch_size is not None:
+            key, *subkeys = jax.random.split(key, 5)
+        else:
+            subkeys = None
+        omega_border = self.sample_in_omega_border_domain(subkeys)
 
-    def _get_omega_operands(self):
+        return key, omega, omega_border
+
+    def _get_omega_operands(
+        self,
+    ) -> tuple[Key, Float[Array, "n dim"], Int, Int, Float[Array, "n"]]:
         return (
-            self._key,
+            self.key,
             self.omega,
             self.curr_omega_idx,
             self.omega_batch_size,
             self.p_omega,
         )
 
-    def inside_batch(self):
+    def inside_batch(
+        self,
+    ) -> tuple["CubicMeshPDEStatio", Float[Array, "omega_batch_size dim"]]:
         r"""
-        Return a batch of points in :math:`\Omega`.
+        Return a batch of points in $\Omega$.
         If all the batches have been seen, we reshuffle them,
         otherwise we just return the next unseen batch.
         """
@@ -670,38 +624,46 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
         bstart = self.curr_omega_idx
         bend = bstart + self.omega_batch_size
 
-        (self._key, self.omega, self.curr_omega_idx) = _reset_or_increment(
-            bend, n_eff, self._get_omega_operands()
+        new_attributes = _reset_or_increment(bend, n_eff, self._get_omega_operands())
+        new = eqx.tree_at(
+            lambda m: (m.key, m.omega, m.curr_omega_idx), self, new_attributes
         )
 
-        # commands below are equivalent to
-        # return self.omega[i:(i+batch_size), 0:dim]
-        return jax.lax.dynamic_slice(
-            self.omega,
-            start_indices=(self.curr_omega_idx, 0),
-            slice_sizes=(self.omega_batch_size, self.dim),
+        return new, jax.lax.dynamic_slice(
+            new.omega,
+            start_indices=(new.curr_omega_idx, 0),
+            slice_sizes=(new.omega_batch_size, new.dim),
         )
 
-    def _get_omega_border_operands(self):
+    def _get_omega_border_operands(
+        self,
+    ) -> tuple[
+        Key, Float[Array, "1 2"] | Float[Array, "(nb//4) 2 4"] | None, Int, Int, None
+    ]:
         return (
-            self._key,
+            self.key,
             self.omega_border,
             self.curr_omega_border_idx,
             self.omega_border_batch_size,
             self.p_border,
         )
 
-    def border_batch(self):
+    def border_batch(
+        self,
+    ) -> tuple[
+        "CubicMeshPDEStatio",
+        Float[Array, "1 1 2"] | Float[Array, "omega_border_batch_size 2 4"] | None,
+    ]:
         r"""
         Return
 
         - The value `None` if `self.omega_border_batch_size` is `None`.
 
-        - a jnp array with two fixed values :math:`(x_{min}, x_{max})` if
+        - a jnp array with two fixed values $(x_{min}, x_{max})$ if
           `self.dim` = 1. There is no sampling here, we return the entire
-          :math:`\partial\Omega`
+          $\partial\Omega$
 
-        - a batch of points in :math:`\partial\Omega` otherwise, stacked by
+        - a batch of points in $\partial\Omega$ otherwise, stacked by
           facet on the last axis.
           If all the batches have been seen, we reshuffle them,
           otherwise we just return the next unseen batch.
@@ -709,236 +671,138 @@ class CubicMeshPDEStatio(DataGeneratorPDEAbstract):
 
         """
         if self.omega_border_batch_size is None:
-            return None
+            return self, None
         if self.dim == 1:
             # 1-D case, no randomness : we always return the whole omega border,
             # i.e. (1, 1, 2) shape jnp.array([[[xmin], [xmax]]]).
-            return self.omega_border[None, None]  # shape is (1, 1, 2)
+            return self, self.omega_border[None, None]  # shape is (1, 1, 2)
         bstart = self.curr_omega_border_idx
         bend = bstart + self.omega_border_batch_size
 
-        (
-            self._key,
-            self.omega_border,
-            self.curr_omega_border_idx,
-        ) = _reset_or_increment(bend, self.nb, self._get_omega_border_operands())
-
-        # commands below are equivalent to
-        # return self.omega[i:(i+batch_size), 0:dim, 0:nb_facets]
-        # and nb_facets = 2 * dimension
-        # but JAX prefer the latter
-        return jax.lax.dynamic_slice(
-            self.omega_border,
-            start_indices=(self.curr_omega_border_idx, 0, 0),
-            slice_sizes=(self.omega_border_batch_size, self.dim, 2 * self.dim),
+        new_attributes = _reset_or_increment(
+            bend, self.nb, self._get_omega_border_operands()
         )
-
-    def get_batch(self):
-        """
-        Generic method to return a batch. Here we call `self.inside_batch()`
-        and `self.border_batch()`
-        """
-        return PDEStatioBatch(
-            inside_batch=self.inside_batch(), border_batch=self.border_batch()
+        new = eqx.tree_at(
+            lambda m: (m.key, m.omega_border, m.curr_omega_border_idx),
+            self,
+            new_attributes,
         )
 
-    def tree_flatten(self):
-        children = (
-            self._key,
-            self.omega,
-            self.omega_border,
-            self.curr_omega_idx,
-            self.curr_omega_border_idx,
-            self.min_pts,
-            self.max_pts,
-            self.p_omega,
-            self.rar_iter_from_last_sampling,
-            self.rar_iter_nb,
+        return new, jax.lax.dynamic_slice(
+            new.omega_border,
+            start_indices=(new.curr_omega_border_idx, 0, 0),
+            slice_sizes=(new.omega_border_batch_size, new.dim, 2 * new.dim),
         )
-        aux_data = {
-            k: vars(self)[k]
-            for k in [
-                "n",
-                "nb",
-                "omega_batch_size",
-                "omega_border_batch_size",
-                "method",
-                "dim",
-                "rar_parameters",
-                "n_start",
-            ]
-        }
-        return (children, aux_data)
 
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
+    def get_batch(self) -> tuple["CubicMeshPDEStatio", PDEStatioBatch]:
         """
-        **Note:** When reconstructing the class, we force ``data_exists=True``
-        in order not to re-generate the data at each flattening and
-        unflattening that happens e.g. during the gradient descent in the
-        optimization process
+        Generic method to return a batch. Here we call `self.inside_batch()`
+        and `self.border_batch()`
         """
-        (
-            key,
-            omega,
-            omega_border,
-            curr_omega_idx,
-            curr_omega_border_idx,
-            min_pts,
-            max_pts,
-            p_omega,
-            rar_iter_from_last_sampling,
-            rar_iter_nb,
-        ) = children
-        # force data_exists=True here in order not to re-generate the data
-        # at each iteration of lax.scan
-        obj = cls(
-            key=key,
-            data_exists=True,
-            min_pts=min_pts,
-            max_pts=max_pts,
-            **aux_data,
-        )
-        obj.omega = omega
-        obj.omega_border = omega_border
-        obj.curr_omega_idx = curr_omega_idx
-        obj.curr_omega_border_idx = curr_omega_border_idx
-        obj.p_omega = p_omega
-        obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
-        obj.rar_iter_nb = rar_iter_nb
-        return obj
+        new, inside_batch = self.inside_batch()
+        new, border_batch = new.border_batch()
+        return new, PDEStatioBatch(inside_batch=inside_batch, border_batch=border_batch)
 
 
-@register_pytree_node_class
 class CubicMeshPDENonStatio(CubicMeshPDEStatio):
-    """
+    r"""
     A class implementing data generator object for non stationary partial
     differential equations. Formally, it extends `CubicMeshPDEStatio`
     to include a temporal batch.
 
-
-    **Note:** CubicMeshPDENonStatio is jittable. Hence it implements the tree_flatten() and
-    tree_unflatten methods.
+    Parameters
+    ----------
+    key : Key
+        Jax random key to sample new time points and to shuffle batches
+    n : Int
+        The number of total $\Omega$ points that will be divided in
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    nb : Int | None
+        The total number of points in $\partial\Omega$.
+        Can be `None` not to lose performance generating the border
+        batch if they are not used
+    nt : Int
+        The number of total time points that will be divided in
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    omega_batch_size : Int
+        The size of the batch of randomly selected points among
+        the `n` points.
+    omega_border_batch_size : Int | None
+        The size of the batch of points randomly selected
+        among the `nb` points.
+        Can be `None` not to lose performance generating the border
+        batch if they are not used
+    temporal_batch_size : Int
+        The size of the batch of randomly selected points among
+        the `nt` points.
+    dim : Int
+        An integer. dimension of $\Omega$ domain
+    min_pts : tuple[tuple[Float, Float], ...]
+        A tuple of minimum values of the domain along each dimension. For a sampling
+        in `n` dimension, this represents $(x_{1, min}, x_{2,min}, ...,
+        x_{n, min})$
+    max_pts : tuple[tuple[Float, Float], ...]
+        A tuple of maximum values of the domain along each dimension. For a sampling
+        in `n` dimension, this represents $(x_{1, max}, x_{2,max}, ...,
+        x_{n,max})$
+    tmin : float
+        The minimum value of the time domain to consider
+    tmax : float
+        The maximum value of the time domain to consider
+    method : str, default="uniform"
+        Either `grid` or `uniform`, default is `uniform`.
+        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 : Dict[str, Int], default=None
+        Default to None: do not use Residual Adaptative Resampling.
+        Otherwise a dictionary with keys. `start_iter`: the iteration at
+        which we start the RAR sampling scheme (we first have a burn in
+        period). `update_every`: the number of gradient steps taken between
+        each appending of collocation points in the RAR algo.
+        `sample_size_omega`: the size of the sample from which we will select new
+        collocation points. `selected_sample_size_omega`: the number of selected
+        points from the sample to be added to the current collocation
+        points.
+    n_start : Int, default=None
+        Defaults to None. The effective size of n used at start time.
+        This value must be
+        provided when rar_parameters is not None. Otherwise we set internally
+        n_start = n and this is hidden from the user.
+        In RAR, n_start
+        then corresponds to the initial number of omega points we train the PINN.
+    nt_start : Int, default=None
+        Defaults to None. A RAR hyper-parameter. Same as ``n_start`` but
+        for times collocation point. See also ``DataGeneratorODE``
+        documentation.
+    cartesian_product : Bool, default=True
+        Defaults to True. Whether we return the cartesian product of the
+        temporal batch with the inside and border batches. If False we just
+        return their concatenation.
     """
 
-    def __init__(
-        self,
-        key,
-        n,
-        nb,
-        nt,
-        omega_batch_size,
-        omega_border_batch_size,
-        temporal_batch_size,
-        dim,
-        min_pts,
-        max_pts,
-        tmin,
-        tmax,
-        method="grid",
-        rar_parameters=None,
-        n_start=None,
-        nt_start=None,
-        cartesian_product=True,
-        data_exists=False,
-    ):
-        r"""
-        Parameters
-        ----------
-        key
-            Jax random key to sample new time points and to shuffle batches
-        n
-            An integer. The number of total :math:`\Omega` points that will be divided in
-            batches. Batches are made so that each data point is seen only
-            once during 1 epoch.
-        nb
-            An integer. The total number of points in :math:`\partial\Omega`.
-            Can be `None` not to lose performance generating the border
-            batch if they are not used
-        nt
-            An integer. The number of total time points that will be divided in
-            batches. Batches are made so that each data point is seen only
-            once during 1 epoch.
-        omega_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `n` points.
-        omega_border_batch_size
-            An integer. The size of the batch of points randomly selected
-            among the `nb` points.
-            Can be `None` not to lose performance generating the border
-            batch if they are not used
-        temporal_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `nt` points.
-        dim
-            An integer. dimension of :math:`\Omega` domain
-        min_pts
-            A tuple of minimum values of the domain along each dimension. For a sampling
-            in `n` dimension, this represents :math:`(x_{1, min}, x_{2,min}, ...,
-            x_{n, min})`
-        max_pts
-            A tuple of maximum values of the domain along each dimension. For a sampling
-            in `n` dimension, this represents :math:`(x_{1, max}, x_{2,max}, ...,
-            x_{n,max})`
-        tmin
-            A float. The minimum value of the time domain to consider
-        tmax
-            A float. The maximum value of the time domain to consider
-        method
-            Either `grid` or `uniform`, default is `grid`.
-            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
-            Default to None: do not use Residual Adaptative Resampling.
-            Otherwise a dictionary with keys. `start_iter`: the iteration at
-            which we start the RAR sampling scheme (we first have a burn in
-            period). `update_every`: the number of gradient steps taken between
-            each appending of collocation points in the RAR algo.
-            `sample_size_omega`: the size of the sample from which we will select new
-            collocation points. `selected_sample_size_omega`: the number of selected
-            points from the sample to be added to the current collocation
-            points.
-        n_start
-            Defaults to None. The effective size of n used at start time.
-            This value must be
-            provided when rar_parameters is not None. Otherwise we set internally
-            n_start = n and this is hidden from the user.
-            In RAR, n_start
-            then corresponds to the initial number of omega points we train the PINN.
-        nt_start
-            Defaults to None. A RAR hyper-parameter. Same as ``n_start`` but
-            for times collocation point. See also ``DataGeneratorODE``
-            documentation.
-        cartesian_product
-            Defaults to True. Whether we return the cartesian product of the
-            temporal batch with the inside and border batches. If False we just
-            return their concatenation.
-        data_exists
-            Must be left to `False` when created by the user. Avoids the
-            regeneration of :math:`\Omega`, :math:`\partial\Omega` and
-            time points at each pytree flattening and unflattening.
+    temporal_batch_size: Int = eqx.field(kw_only=True)
+    tmin: Float = eqx.field(kw_only=True)
+    tmax: Float = eqx.field(kw_only=True)
+    nt: Int = eqx.field(kw_only=True)
+    temporal_batch_size: Int = eqx.field(kw_only=True, static=True)
+    cartesian_product: Bool = eqx.field(kw_only=True, default=True, static=True)
+    nt_start: int = eqx.field(kw_only=True, default=None, static=True)
+
+    p_times: Array = eqx.field(init=False)
+    curr_time_idx: Int = eqx.field(init=False)
+    times: Array = eqx.field(init=False)
+
+    def __post_init__(self):
         """
-        super().__init__(
-            key,
-            n,
-            nb,
-            omega_batch_size,
-            omega_border_batch_size,
-            dim,
-            min_pts,
-            max_pts,
-            method,
-            rar_parameters,
-            n_start,
-            data_exists,
-        )
-        self.temporal_batch_size = temporal_batch_size
-        self.tmin = tmin
-        self.tmax = tmax
-        self.nt = nt
+        Note that neither __init__ or __post_init__ are called when udating a
+        Module with eqx.tree_at!
+        """
+        super().__post_init__()  # because __init__ or __post_init__ of Base
+        # class is not automatically called
 
-        self.cartesian_product = cartesian_product
         if not self.cartesian_product:
             if self.temporal_batch_size != self.omega_batch_size:
                 raise ValueError(
@@ -968,46 +832,54 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
             self.p_times,
             _,
             _,
-        ) = _check_and_set_rar_parameters(rar_parameters, n=nt, n_start=nt_start)
-
-        if not self.data_exists:
-            # Useful when using a lax.scan with pytree
-            # Optionally can tell JAX not to re-generate data
-            self.curr_time_idx = 0
-            self.generate_data_nonstatio()
-            self._key, self.times, _ = _reset_batch_idx_and_permute(
-                self._get_time_operands()
-            )
-
-    def sample_in_time_domain(self, n_samples):
-        self._key, subkey = random.split(self._key, 2)
-        return random.uniform(subkey, (n_samples,), minval=self.tmin, maxval=self.tmax)
+        ) = _check_and_set_rar_parameters(self.rar_parameters, self.nt, self.nt_start)
+
+        self.curr_time_idx = jnp.iinfo(jnp.int32).max - self.temporal_batch_size - 1
+        self.key, _ = jax.random.split(self.key, 2)  # to make it equivalent to
+        # the call to _reset_batch_idx_and_permute in legacy DG
+        self.key, self.times = self.generate_time_data(self.key)
+        # see explaination in DataGeneratorODE for the key
+
+    def sample_in_time_domain(
+        self, key: Key, sample_size: Int = None
+    ) -> Float[Array, "nt"]:
+        return jax.random.uniform(
+            key,
+            (self.nt if sample_size is None else sample_size,),
+            minval=self.tmin,
+            maxval=self.tmax,
+        )
 
-    def _get_time_operands(self):
+    def _get_time_operands(
+        self,
+    ) -> tuple[Key, Float[Array, "nt"], Int, Int, Float[Array, "nt"]]:
         return (
-            self._key,
+            self.key,
             self.times,
             self.curr_time_idx,
             self.temporal_batch_size,
             self.p_times,
         )
 
-    def generate_data_nonstatio(self):
-        r"""
+    def generate_time_data(self, key: Key) -> tuple[Key, Float[Array, "nt"]]:
+        """
         Construct a complete set of `self.nt` time points according to the
-        specified `self.method`. This completes the `super` function
-        `generate_data()` which generates :math:`\Omega` and
-        :math:`\partial\Omega` points.
+        specified `self.method`
+
+        Note that self.times has always size self.nt and not self.nt_start, even
+        in RAR scheme, we must allocate all the collocation points
         """
+        key, subkey = jax.random.split(key, 2)
         if self.method == "grid":
-            self.partial_times = (self.tmax - self.tmin) / self.nt
-            self.times = jnp.arange(self.tmin, self.tmax, self.partial_times)
-        elif self.method == "uniform":
-            self.times = self.sample_in_time_domain(self.nt)
-        else:
-            raise ValueError("Method " + self.method + " is not implemented.")
+            partial_times = (self.tmax - self.tmin) / self.nt
+            return key, jnp.arange(self.tmin, self.tmax, partial_times)
+        if self.method == "uniform":
+            return key, self.sample_in_time_domain(subkey)
+        raise ValueError("Method " + self.method + " is not implemented.")
 
-    def temporal_batch(self):
+    def temporal_batch(
+        self,
+    ) -> tuple["CubicMeshPDENonStatio", Float[Array, "temporal_batch_size"]]:
         """
         Return a batch of time points. If all the batches have been seen, we
         reshuffle them, otherwise we just return the next unseen batch.
@@ -1018,254 +890,344 @@ class CubicMeshPDENonStatio(CubicMeshPDEStatio):
         # Compute the effective number of used collocation points
         if self.rar_parameters is not None:
             nt_eff = (
-                self.n_start
+                self.nt_start
                 + self.rar_iter_nb * self.rar_parameters["selected_sample_size_times"]
             )
         else:
             nt_eff = self.nt
 
-        (self._key, self.times, self.curr_time_idx) = _reset_or_increment(
-            bend, nt_eff, self._get_time_operands()
+        new_attributes = _reset_or_increment(bend, nt_eff, self._get_time_operands())
+        new = eqx.tree_at(
+            lambda m: (m.key, m.times, m.curr_time_idx), self, new_attributes
         )
 
-        # commands below are equivalent to
-        # return self.times[i:(i+t_batch_size)]
-        # but JAX prefer the latter
-        return jax.lax.dynamic_slice(
-            self.times,
-            start_indices=(self.curr_time_idx,),
-            slice_sizes=(self.temporal_batch_size,),
+        return new, jax.lax.dynamic_slice(
+            new.times,
+            start_indices=(new.curr_time_idx,),
+            slice_sizes=(new.temporal_batch_size,),
         )
 
-    def get_batch(self):
+    def get_batch(self) -> tuple["CubicMeshPDENonStatio", PDENonStatioBatch]:
         """
         Generic method to return a batch. Here we call `self.inside_batch()`,
         `self.border_batch()` and `self.temporal_batch()`
         """
-        x = self.inside_batch()
-        dx = self.border_batch()
-        t = self.temporal_batch().reshape(self.temporal_batch_size, 1)
+        new, x = self.inside_batch()
+        new, dx = new.border_batch()
+        new, t = new.temporal_batch()
+        t = t.reshape(new.temporal_batch_size, 1)
 
-        if self.cartesian_product:
+        if new.cartesian_product:
             t_x = make_cartesian_product(t, x)
         else:
             t_x = jnp.concatenate([t, x], axis=1)
 
         if dx is not None:
-            t_ = t.reshape(self.temporal_batch_size, 1, 1)
+            t_ = t.reshape(new.temporal_batch_size, 1, 1)
             t_ = jnp.repeat(t_, dx.shape[-1], axis=2)
-            if self.cartesian_product or self.dim == 1:
+            if new.cartesian_product or new.dim == 1:
                 t_dx = make_cartesian_product(t_, dx)
             else:
                 t_dx = jnp.concatenate([t_, dx], axis=1)
         else:
             t_dx = None
 
-        return PDENonStatioBatch(times_x_inside_batch=t_x, times_x_border_batch=t_dx)
+        return new, PDENonStatioBatch(
+            times_x_inside_batch=t_x, times_x_border_batch=t_dx
+        )
 
-    def tree_flatten(self):
-        children = (
-            self._key,
-            self.omega,
-            self.omega_border,
-            self.times,
-            self.curr_omega_idx,
-            self.curr_omega_border_idx,
-            self.curr_time_idx,
-            self.min_pts,
-            self.max_pts,
-            self.tmin,
-            self.tmax,
-            self.p_times,
-            self.p_omega,
-            self.rar_iter_from_last_sampling,
-            self.rar_iter_nb,
+
+class DataGeneratorObservations(eqx.Module):
+    r"""
+    Despite the class name, it is rather a dataloader from user provided
+    observations that will be used for the observations loss
+
+    Parameters
+    ----------
+    key : Key
+        Jax random key to shuffle batches
+    obs_batch_size : Int
+        The size of the batch of randomly selected points among
+        the `n` points. `obs_batch_size` will be the same for all
+        elements of the return observation dict batch.
+        NOTE: no check is done BUT users should be careful that
+        `obs_batch_size` must be equal to `temporal_batch_size` or
+        `omega_batch_size` or the product of both. In the first case, the
+        present DataGeneratorObservations instance complements an ODEBatch,
+        PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+        = False). In the second case, `obs_batch_size` =
+        `temporal_batch_size * omega_batch_size` if the present
+        DataGeneratorParameter complements a PDENonStatioBatch
+        with self.cartesian_product = True
+    observed_pinn_in : Float[Array, "n_obs nb_pinn_in"]
+        Observed values corresponding to the input of the PINN
+        (eg. the time at which we recorded the observations). The first
+        dimension must corresponds to the number of observed_values.
+        The second dimension depends on the input dimension of the PINN,
+        that is `1` for ODE, `n_dim_x` for stationnary PDE and `n_dim_x + 1`
+        for non-stationnary PDE.
+    observed_values : Float[Array, "n_obs, nb_pinn_out"]
+        Observed values that the PINN should learn to fit. The first
+        dimension must be aligned with observed_pinn_in.
+    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
+        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.
+        Optional argument.
+    sharding_device : jax.sharding.Sharding, default=None
+        Default None. An optional sharding object to constraint the storage
+        of observed inputs, values and parameters. Typically, a
+        SingleDeviceSharding(cpu_device) to avoid loading on GPU huge
+        datasets of observations. Note that computations for **batches**
+        can still be performed on other devices (*e.g.* GPU, TPU or
+        any pre-defined Sharding) thanks to the `obs_batch_sharding`
+        arguments of `jinns.solve()`. Read the docs for more info.
+    """
+
+    key: Key
+    obs_batch_size: Int = 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)
+
+    n: Int = eqx.field(init=False)
+    curr_idx: Int = eqx.field(init=False)
+    indices: Array = eqx.field(init=False)
+
+    def __post_init__(self):
+        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 len(self.observed_pinn_in.shape) > 2:
+            raise ValueError("self.observed_pinn_in must have 2 dimensions")
+        if len(self.observed_values.shape) == 1:
+            self.observed_values = self.observed_values[:, None]
+        if len(self.observed_values.shape) > 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"
+                )
+
+        self.n = self.observed_pinn_in.shape[0]
+
+        if self.sharding_device is not None:
+            self.observed_pinn_in = jax.lax.with_sharding_constraint(
+                self.observed_pinn_in, self.sharding_device
+            )
+            self.observed_values = jax.lax.with_sharding_constraint(
+                self.observed_values, self.sharding_device
+            )
+            self.observed_eq_params = jax.lax.with_sharding_constraint(
+                self.observed_eq_params, self.sharding_device
+            )
+
+        self.curr_idx = jnp.iinfo(jnp.int32).max - self.obs_batch_size - 1
+        # For speed and to avoid duplicating data what is really
+        # shuffled is a vector of indices
+        if self.sharding_device is not None:
+            self.indices = jax.lax.with_sharding_constraint(
+                jnp.arange(self.n), self.sharding_device
+            )
+        else:
+            self.indices = jnp.arange(self.n)
+
+        # recall post_init is the only place with _init_ where we can set
+        # self attribute in a in-place way
+        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]:
+        return (
+            self.key,
+            self.indices,
+            self.curr_idx,
+            self.obs_batch_size,
+            None,
         )
-        aux_data = {
-            k: vars(self)[k]
-            for k in [
-                "n",
-                "nb",
-                "nt",
-                "omega_batch_size",
-                "omega_border_batch_size",
-                "temporal_batch_size",
-                "method",
-                "dim",
-                "rar_parameters",
-                "n_start",
-                "nt_start",
-                "cartesian_product",
-            ]
-        }
-        return (children, aux_data)
 
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
+    def obs_batch(
+        self,
+    ) -> tuple[
+        "DataGeneratorObservations", Dict[str, Float[Array, "obs_batch_size dim"]]
+    ]:
         """
-        **Note:** When reconstructing the class, we force ``data_exists=True``
-        in order not to re-generate the data at each flattening and
-        unflattening that happens e.g. during the gradient descent in the
-        optimization process
+        Return a dictionary with (keys, values): (pinn_in, a mini batch of pinn
+        inputs), (obs, a mini batch of corresponding observations), (eq_params,
+        a dictionary with entry names found in `params["eq_params"]` and values
+        giving the correspond parameter value for the couple
+        (input, observation) mentioned before).
+        It can also be a dictionary of dictionaries as described above if
+        observed_pinn_in, observed_values, etc. are dictionaries with keys
+        representing the PINNs.
         """
-        (
-            key,
-            omega,
-            omega_border,
-            times,
-            curr_omega_idx,
-            curr_omega_border_idx,
-            curr_time_idx,
-            min_pts,
-            max_pts,
-            tmin,
-            tmax,
-            p_times,
-            p_omega,
-            rar_iter_from_last_sampling,
-            rar_iter_nb,
-        ) = children
-        obj = cls(
-            key=key,
-            data_exists=True,
-            min_pts=min_pts,
-            max_pts=max_pts,
-            tmin=tmin,
-            tmax=tmax,
-            **aux_data,
+
+        new_attributes = _reset_or_increment(
+            self.curr_idx + self.obs_batch_size, self.n, self._get_operands()
+        )
+        new = eqx.tree_at(
+            lambda m: (m.key, m.indices, m.curr_idx), self, new_attributes
+        )
+
+        minib_indices = jax.lax.dynamic_slice(
+            new.indices,
+            start_indices=(new.curr_idx,),
+            slice_sizes=(new.obs_batch_size,),
         )
-        obj.omega = omega
-        obj.omega_border = omega_border
-        obj.times = times
-        obj.curr_omega_idx = curr_omega_idx
-        obj.curr_omega_border_idx = curr_omega_border_idx
-        obj.curr_time_idx = curr_time_idx
-        obj.p_times = p_times
-        obj.p_omega = p_omega
-        obj.rar_iter_from_last_sampling = rar_iter_from_last_sampling
-        obj.rar_iter_nb = rar_iter_nb
-
-        return obj
-
-
-@register_pytree_node_class
-class DataGeneratorParameter:
-    """
-    A data generator for additional unidimensional parameter(s)
-    """
 
-    def __init__(
+        obs_batch = {
+            "pinn_in": jnp.take(
+                new.observed_pinn_in, minib_indices, unique_indices=True, axis=0
+            ),
+            "val": jnp.take(
+                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),
+                new.observed_eq_params,
+            ),
+        }
+        return new, obs_batch
+
+    def get_batch(
         self,
-        key,
-        n,
-        param_batch_size,
-        param_ranges=None,
-        method="grid",
-        user_data=None,
-        data_exists=False,
-    ):
-        r"""
-        Parameters
-        ----------
-        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
-        n
-            An integer. The number of total points that will be divided in
-            batches. Batches are made so that each data point is seen only
-            once during 1 epoch.
-        param_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `n` points. `param_batch_size` will be the same for all
-            additional batch of parameter.
-            NOTE: no check is done BUT users should be careful that
-            `param_batch_size` must be equal to `temporal_batch_size` or
-            `omega_batch_size` or the product of both. In the first case, the
-            present DataGeneratorParameter instance complements an ODEBatch, a
-            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
-            = False). In the second case, `param_batch_size` =
-            `temporal_batch_size * omega_batch_size` if the present
-            DataGeneratorParameter complements a PDENonStatioBatch
-            with self.cartesian_product = True
-        param_ranges
-            A dict. A dict of tuples (min, max), which
-            reprensents the range of real numbers where to sample batches (of
-            length `param_batch_size` among `n` points).
-            The key corresponds to the parameter name. The keys must match the
-            keys in `params["eq_params"]`.
-            By providing several entries in this dictionary we can sample
-            an arbitrary number of parameters.
-            __Note__ that we currently only support unidimensional parameters
-        method
-            Either `grid` or `uniform`, default is `grid`. `grid` means
-            regularly spaced points over the domain. `uniform` means uniformly
-            sampled points over the domain
-        data_exists
-            Must be left to `False` when created by the user. Avoids the
-            regeneration of :math:`\Omega`, :math:`\partial\Omega` and
-            time points at each pytree flattening and unflattening.
-        user_data
-            A dictionary containing user-provided data for parameters.
-            As for `param_ranges`, the key corresponds to the parameter name,
-            the keys must match the keys in `params["eq_params"]` and only
-            unidimensional arrays are supported. Therefore, the jnp arrays
-            found at `user_data[k]` must have shape `(n, 1)` or `(n,)`.
-            Note that if the same key appears in `param_ranges` and `user_data`
-            priority goes for the content in `user_data`.
-            Defaults to None.
+    ) -> tuple[
+        "DataGeneratorObservations", Dict[str, Float[Array, "obs_batch_size dim"]]
+    ]:
         """
-        self.data_exists = data_exists
-        self.method = method
+        Generic method to return a batch
+        """
+        return self.obs_batch()
+
+
+class DataGeneratorParameter(eqx.Module):
+    r"""
+    A data generator for additional unidimensional parameter(s)
+
+    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
+    n : Int
+        The number of total points that will be divided in
+        batches. Batches are made so that each data point is seen only
+        once during 1 epoch.
+    param_batch_size : Int
+        The size of the batch of randomly selected points among
+        the `n` points. `param_batch_size` will be the same for all
+        additional batch of parameter.
+        NOTE: no check is done BUT users should be careful that
+        `param_batch_size` must be equal to `temporal_batch_size` or
+        `omega_batch_size` or the product of both. In the first case, the
+        present DataGeneratorParameter instance complements an ODEBatch, a
+        PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+        = False). In the second case, `param_batch_size` =
+        `temporal_batch_size * omega_batch_size` if the present
+        DataGeneratorParameter complements a PDENonStatioBatch
+        with self.cartesian_product = True
+    param_ranges : Dict[str, tuple[Float, Float] | None, default={}
+        A dict. A dict of tuples (min, max), which
+        reprensents the range of real numbers where to sample batches (of
+        length `param_batch_size` among `n` points).
+        The key corresponds to the parameter name. The keys must match the
+        keys in `params["eq_params"]`.
+        By providing several entries in this dictionary we can sample
+        an arbitrary number of parameters.
+        **Note** that we currently only support unidimensional parameters.
+        This argument can be done if we only use `user_data`.
+    method : str, default="uniform"
+        Either `grid` or `uniform`, default is `uniform`. `grid` means
+        regularly spaced points over the domain. `uniform` means uniformly
+        sampled points over the domain
+    user_data : Dict[str, Float[Array, "n"]] | None, default={}
+        A dictionary containing user-provided data for parameters.
+        As for `param_ranges`, the key corresponds to the parameter name,
+        the keys must match the keys in `params["eq_params"]` and only
+        unidimensional arrays are supported. Therefore, the jnp arrays
+        found at `user_data[k]` must have shape `(n, 1)` or `(n,)`.
+        Note that if the same key appears in `param_ranges` and `user_data`
+        priority goes for the content in `user_data`.
+        Defaults to None.
+    """
+
+    keys: Key | Dict[str, Key]
+    n: Int
+    param_batch_size: Int = eqx.field(static=True)
+    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(
+        static=True, default_factory=lambda: {}
+    )
+
+    curr_param_idx: Dict[str, Int] = eqx.field(init=False)
+    param_n_samples: Dict[str, Array] = eqx.field(init=False)
 
-        if n < param_batch_size:
+    def __post_init__(self):
+        if self.user_data is None:
+            self.user_data = {}
+        if self.param_ranges is None:
+            self.param_ranges = {}
+        if self.n < self.param_batch_size:
             raise ValueError(
-                f"Number of data points ({n}) is smaller than the"
-                f"number of batch points ({param_batch_size})."
+                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))))
+
+        self.curr_param_idx = {}
+        for k in self.keys.keys():
+            self.curr_param_idx[k] = (
+                jnp.iinfo(jnp.int32).max - self.param_batch_size - 1
             )
 
-        if user_data is None:
-            user_data = {}
-        if param_ranges is None:
-            param_ranges = {}
-        if not isinstance(key, dict):
-            all_keys = set().union(param_ranges, user_data)
-            self._keys = dict(zip(all_keys, jax.random.split(key, len(all_keys))))
-        else:
-            self._keys = key
-        self.n = n
-        self.param_batch_size = param_batch_size
-        self.param_ranges = param_ranges
-        self.user_data = user_data
-
-        if not self.data_exists:
-            self.generate_data()
-            # The previous call to self.generate_data() has created
-            # the dict self.param_n_samples and then we will only use this one
-            # because it has merged the scattered data between `user_data` and
-            # `param_ranges`
-            self.curr_param_idx = {}
-            for k in self.param_n_samples.keys():
-                self.curr_param_idx[k] = 0
-                (
-                    self._keys[k],
-                    self.param_n_samples[k],
-                    _,
-                ) = _reset_batch_idx_and_permute(self._get_param_operands(k))
-
-    def generate_data(self):
+        # 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)
+
+    def generate_data(
+        self, keys: Dict[str, Key]
+    ) -> tuple[Dict[str, Key], Dict[str, Float[Array, "n"]]]:
         """
         Generate parameter samples, either through generation
         or using user-provided data.
         """
-        self.param_n_samples = {}
+        param_n_samples = {}
 
         all_keys = set().union(self.param_ranges, self.user_data)
         for k in all_keys:
-            if self.user_data and k in self.user_data:
+            if (
+                self.user_data
+                and k in self.user_data.keys()  # pylint: disable=no-member
+            ):
                 if self.user_data[k].shape == (self.n, 1):
-                    self.param_n_samples[k] = self.user_data[k]
+                    param_n_samples[k] = self.user_data[k]
                 if self.user_data[k].shape == (self.n,):
-                    self.param_n_samples[k] = self.user_data[k][:, None]
+                    param_n_samples[k] = self.user_data[k][:, None]
                 else:
                     raise ValueError(
                         "Wrong shape for user provided parameters"
@@ -1274,23 +1236,25 @@ class DataGeneratorParameter:
             else:
                 if self.method == "grid":
                     xmin, xmax = self.param_ranges[k][0], self.param_ranges[k][1]
-                    self.partial = (xmax - xmin) / self.n
+                    partial = (xmax - xmin) / self.n
                     # shape (n, 1)
-                    self.param_n_samples[k] = jnp.arange(xmin, xmax, self.partial)[
-                        :, None
-                    ]
+                    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]
-                    self._keys[k], subkey = random.split(self._keys[k], 2)
-                    self.param_n_samples[k] = random.uniform(
+                    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
                     )
                 else:
                     raise ValueError("Method " + self.method + " is not implemented.")
 
-    def _get_param_operands(self, k):
+        return keys, param_n_samples
+
+    def _get_param_operands(
+        self, k: str
+    ) -> tuple[Key, Float[Array, "n"], Int, Int, None]:
         return (
-            self._keys[k],
+            self.keys[k],
             self.param_n_samples[k],
             self.curr_param_idx[k],
             self.param_batch_size,
@@ -1315,26 +1279,28 @@ class DataGeneratorParameter:
             _reset_or_increment_wrapper,
             self.param_n_samples,
             self.curr_param_idx,
-            self._keys,
+            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
-        (
-            self._keys,
-            self.param_n_samples,
-            self.curr_param_idx,
-        ) = jax.tree_util.tree_transpose(
-            jax.tree_util.tree_structure(self._keys),
+        new_attributes = jax.tree_util.tree_transpose(
+            jax.tree_util.tree_structure(self.keys),
             jax.tree_util.tree_structure([0, 0, 0]),
             res,
         )
 
-        return jax.tree_util.tree_map(
+        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=(self.param_batch_size, 1)
+                p, start_indices=(q, 0), slice_sizes=(new.param_batch_size, 1)
             ),
-            self.param_n_samples,
-            self.curr_param_idx,
+            new.param_n_samples,
+            new.curr_param_idx,
         )
 
     def get_batch(self):
@@ -1343,251 +1309,9 @@ class DataGeneratorParameter:
         """
         return self.param_batch()
 
-    def tree_flatten(self):
-        children = (
-            self._keys,
-            self.param_n_samples,
-            self.curr_param_idx,
-        )
-        aux_data = {
-            k: vars(self)[k]
-            for k in ["n", "param_batch_size", "method", "param_ranges", "user_data"]
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (
-            keys,
-            param_n_samples,
-            curr_param_idx,
-        ) = children
-        obj = cls(
-            key=keys,
-            data_exists=True,
-            **aux_data,
-        )
-        obj.param_n_samples = param_n_samples
-        obj.curr_param_idx = curr_param_idx
-        return obj
-
-
-@register_pytree_node_class
-class DataGeneratorObservations:
-    """
-    Despite the class name, it is rather a dataloader from user provided
-    observations that will be used for the observations loss
-    """
-
-    def __init__(
-        self,
-        key,
-        obs_batch_size,
-        observed_pinn_in,
-        observed_values,
-        observed_eq_params=None,
-        data_exists=False,
-        sharding_device=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        key
-            Jax random key to sample new time points and to shuffle batches
-        obs_batch_size
-            An integer. The size of the batch of randomly selected points among
-            the `n` points. `obs_batch_size` will be the same for all
-            elements of the obs dict.
-            NOTE: no check is done BUT users should be careful that
-            `obs_batch_size` must be equal to `temporal_batch_size` or
-            `omega_batch_size` or the product of both. In the first case, the
-            present DataGeneratorObservations instance complements an ODEBatch,
-            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
-            = False). In the second case, `obs_batch_size` =
-            `temporal_batch_size * omega_batch_size` if the present
-            DataGeneratorParameter complements a PDENonStatioBatch
-            with self.cartesian_product = True
-        observed_pinn_in
-            A jnp.array with 2 dimensions.
-            Observed values corresponding to the input of the PINN
-            (eg. the time at which we recorded the observations). The first
-            dimension must corresponds to the number of observed_values and
-            observed_eq_params. The second dimension depends on the input dimension of the PINN, that is `1` for ODE, `n_dim_x` for stationnary PDE and `n_dim_x + 1` for non-stationnary PDE.
-        observed_values
-            A jnp.array with 2 dimensions.
-            Observed values that the PINN should learn to fit. The first dimension must be aligned with observed_pinn_in and the values of observed_eq_params.
-        observed_eq_params
-            Optional. Default is None. 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
-            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.
-        data_exists
-            Must be left to `False` when created by the user. Avoids the
-            resetting of curr_idx at each pytree flattening and unflattening.
-        sharding_device
-            Default None. An optional sharding object to constraint the storage
-            of observed inputs, values and parameters. Typically, a
-            SingleDeviceSharding(cpu_device) to avoid loading on GPU huge
-            datasets of observations. Note that computations for **batches**
-            can still be performed on other devices (*e.g.* GPU, TPU or
-            any pre-defined Sharding) thanks to the `obs_batch_sharding`
-            arguments of `jinns.solve()`. Read the docs for more info.
-
-        """
-        if observed_eq_params is None:
-            observed_eq_params = {}
-
-        if not data_exists:
-            self.observed_eq_params = observed_eq_params.copy()
-        else:
-            # avoid copying when in flatten/unflatten
-            self.observed_eq_params = observed_eq_params
-
-        if observed_pinn_in.shape[0] != observed_values.shape[0]:
-            raise ValueError(
-                "observed_pinn_in and observed_values must have same first axis"
-            )
-        for _, v in self.observed_eq_params.items():
-            if v.shape[0] != observed_pinn_in.shape[0]:
-                raise ValueError(
-                    "observed_pinn_in and the values of"
-                    " observed_eq_params must have the same first axis"
-                )
-        if len(observed_pinn_in.shape) == 1:
-            observed_pinn_in = observed_pinn_in[:, None]
-        if len(observed_pinn_in.shape) > 2:
-            raise ValueError("observed_pinn_in must have 2 dimensions")
-        if len(observed_values.shape) == 1:
-            observed_values = observed_values[:, None]
-        if len(observed_values.shape) > 2:
-            raise ValueError("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"
-                )
-
-        self.n = observed_pinn_in.shape[0]
-        self._key = key
-        self.obs_batch_size = obs_batch_size
-
-        self.data_exists = data_exists
-        if not self.data_exists and sharding_device is not None:
-            self.observed_pinn_in = jax.lax.with_sharding_constraint(
-                observed_pinn_in, sharding_device
-            )
-            self.observed_values = jax.lax.with_sharding_constraint(
-                observed_values, sharding_device
-            )
-            self.observed_eq_params = jax.lax.with_sharding_constraint(
-                self.observed_eq_params, sharding_device
-            )
-        else:
-            self.observed_pinn_in = observed_pinn_in
-            self.observed_values = observed_values
-
-        if not self.data_exists:
-            self.curr_idx = 0
-            # NOTE for speed and to avoid duplicating data what is really
-            # shuffled is a vector of indices
-            indices = jnp.arange(self.n)
-            if sharding_device is not None:
-                self.indices = jax.lax.with_sharding_constraint(
-                    indices, sharding_device
-                )
-            else:
-                self.indices = indices
-            self._key, self.indices, _ = _reset_batch_idx_and_permute(
-                self._get_operands()
-            )
-
-    def _get_operands(self):
-        return (
-            self._key,
-            self.indices,
-            self.curr_idx,
-            self.obs_batch_size,
-            None,
-        )
 
-    def obs_batch(self):
-        """
-        Return a dictionary with (keys, values): (pinn_in, a mini batch of pinn
-        inputs), (obs, a mini batch of corresponding observations), (eq_params,
-        a dictionary with entry names found in `params["eq_params"]` and values
-        giving the correspond parameter value for the couple
-        (input, observation) mentioned before).
-        It can also be a dictionary of dictionaries as described above if
-        observed_pinn_in, observed_values, etc. are dictionaries with keys
-        representing the PINNs.
-        """
-
-        (self._key, self.indices, self.curr_idx) = _reset_or_increment(
-            self.curr_idx + self.obs_batch_size, self.n, self._get_operands()
-        )
-
-        minib_indices = jax.lax.dynamic_slice(
-            self.indices,
-            start_indices=(self.curr_idx,),
-            slice_sizes=(self.obs_batch_size,),
-        )
-
-        obs_batch = {
-            "pinn_in": jnp.take(
-                self.observed_pinn_in, minib_indices, unique_indices=True, axis=0
-            ),
-            "val": jnp.take(
-                self.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),
-                self.observed_eq_params,
-            ),
-        }
-        return obs_batch
-
-    def get_batch(self):
-        """
-        Generic method to return a batch
-        """
-        return self.obs_batch()
-
-    def tree_flatten(self):
-        children = (self._key, self.curr_idx, self.indices)
-        aux_data = {
-            k: vars(self)[k]
-            for k in [
-                "obs_batch_size",
-                "observed_pinn_in",
-                "observed_values",
-                "observed_eq_params",
-            ]
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (key, curr_idx, indices) = children
-        obj = cls(
-            key=key,
-            data_exists=True,
-            obs_batch_size=aux_data["obs_batch_size"],
-            observed_pinn_in=aux_data["observed_pinn_in"],
-            observed_values=aux_data["observed_values"],
-            observed_eq_params=aux_data["observed_eq_params"],
-        )
-        obj.curr_idx = curr_idx
-        obj.indices = indices
-        return obj
-
-
-@register_pytree_node_class
-class DataGeneratorObservationsMultiPINNs:
-    """
+class DataGeneratorObservationsMultiPINNs(eqx.Module):
+    r"""
     Despite the class name, it is rather a dataloader from user provided
     observations that will be used for the observations loss.
     This is the DataGenerator to use when dealing with multiple PINNs
@@ -1597,146 +1321,123 @@ class DataGeneratorObservationsMultiPINNs:
     applied in `constraints_system_loss_apply` and in this case the
     batch.obs_batch_dict is a dict of obs_batch_dict over which the tree_map
     applies (we select the obs_batch_dict corresponding to its `u_dict` entry)
+
+    Parameters
+    ----------
+    obs_batch_size : Int
+        The size of the batch of randomly selected observations
+        `obs_batch_size` will be the same for all the
+        elements of the obs dict.
+        NOTE: no check is done BUT users should be careful that
+        `obs_batch_size` must be equal to `temporal_batch_size` or
+        `omega_batch_size` or the product of both. In the first case, the
+        present DataGeneratorObservations instance complements an ODEBatch,
+        PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
+        = False). In the second case, `obs_batch_size` =
+        `temporal_batch_size * omega_batch_size` if the present
+        DataGeneratorParameter complements a PDENonStatioBatch
+        with self.cartesian_product = True
+    observed_pinn_in_dict : Dict[str, Float[Array, "n_obs nb_pinn_in"] | None]
+        A dict of observed_pinn_in as defined in DataGeneratorObservations.
+        Keys must be that of `u_dict`.
+        If no observation exists for a particular entry of `u_dict` the
+        corresponding key must still exist in observed_pinn_in_dict with
+        value None
+    observed_values_dict : Dict[str, Float[Array, "n_obs, nb_pinn_out"] | None]
+        A dict of observed_values as defined in DataGeneratorObservations.
+        Keys must be that of `u_dict`.
+        If no observation exists for a particular entry of `u_dict` the
+        corresponding key must still exist in observed_values_dict with
+        value None
+    observed_eq_params_dict : Dict[str, Dict[str, Float[Array, "n_obs 1"]]]
+        A dict of observed_eq_params as defined in DataGeneratorObservations.
+        Keys must be that of `u_dict`.
+        **Note**: if no observation exists for a particular entry of `u_dict` the
+        corresponding key must still exist in observed_eq_params_dict with
+        value `{}` (empty dictionnary).
+    key
+        Jax random key to shuffle batches.
     """
 
-    def __init__(
-        self,
-        obs_batch_size,
-        observed_pinn_in_dict,
-        observed_values_dict,
-        observed_eq_params_dict=None,
-        data_gen_obs_exists=False,
-        key=None,
-    ):
-        r"""
-        Parameters
-        ----------
-        obs_batch_size
-            An integer. The size of the batch of randomly selected observations
-            `obs_batch_size` will be the same for all the
-            elements of the obs dict.
-            NOTE: no check is done BUT users should be careful that
-            `obs_batch_size` must be equal to `temporal_batch_size` or
-            `omega_batch_size` or the product of both. In the first case, the
-            present DataGeneratorObservations instance complements an ODEBatch,
-            PDEStatioBatch or a PDENonStatioBatch (with self.cartesian_product
-            = False). In the second case, `obs_batch_size` =
-            `temporal_batch_size * omega_batch_size` if the present
-            DataGeneratorParameter complements a PDENonStatioBatch
-            with self.cartesian_product = True
-        observed_pinn_in_dict
-            A dict of observed_pinn_in as defined in DataGeneratorObservations.
-            Keys must be that of `u_dict`.
-            If no observation exists for a particular entry of `u_dict` the
-            corresponding key must still exist in observed_pinn_in_dict with
-            value None
-        observed_values_dict
-            A dict of observed_values as defined in DataGeneratorObservations.
-            Keys must be that of `u_dict`.
-            If no observation exists for a particular entry of `u_dict` the
-            corresponding key must still exist in observed_values_dict with
-            value None
-        observed_eq_params_dict
-            A dict of observed_eq_params as defined in DataGeneratorObservations.
-            Keys must be that of `u_dict`.
-            If no observation exists for a particular entry of `u_dict` the
-            corresponding key must still exist in observed_eq_params_dict with
-            value None
-        data_gen_obs_exists
-            Must be left to `False` when created by the user. Avoids the
-            regeneration the subclasses DataGeneratorObservations
-            at each pytree flattening and unflattening.
-        key
-            Jax random key to sample new time points and to shuffle batches.
-            Optional if data_gen_obs_exists is True
-        """
-        if (
-            observed_pinn_in_dict is None or observed_values_dict is None
-        ) and not data_gen_obs_exists:
+    obs_batch_size: Int
+    observed_pinn_in_dict: Dict[str, Float[Array, "n_obs nb_pinn_in"] | None]
+    observed_values_dict: Dict[str, Float[Array, "n_obs nb_pinn_out"] | None]
+    observed_eq_params_dict: Dict[str, Dict[str, Float[Array, "n_obs 1"]]] = eqx.field(
+        default=None, kw_only=True
+    )
+    key: InitVar[Key]
+
+    data_gen_obs: Dict[str, "DataGeneratorObservations"] = eqx.field(init=False)
+
+    def __post_init__(self, key):
+        if self.observed_pinn_in_dict is None or self.observed_values_dict is None:
+            raise ValueError(
+                "observed_pinn_in_dict and observed_values_dict " "must be provided"
+            )
+        if self.observed_pinn_in_dict.keys() != self.observed_values_dict.keys():
             raise ValueError(
-                "observed_pinn_in_dict and observed_values_dict "
-                "must be provided with data_gen_obs_exists is False"
+                "Keys must be the same in observed_pinn_in_dict"
+                " and observed_values_dict"
             )
-        self.obs_batch_size = obs_batch_size
-        self.data_gen_obs_exists = data_gen_obs_exists
 
-        if not self.data_gen_obs_exists:
-            if observed_pinn_in_dict.keys() != observed_values_dict.keys():
-                raise ValueError(
-                    "Keys must be the same in observed_pinn_in_dict"
-                    " and observed_values_dict"
-                )
-            if (
-                observed_eq_params_dict is not None
-                and observed_pinn_in_dict.keys() != observed_eq_params_dict.keys()
-            ):
-                raise ValueError(
-                    "Keys must be the same in observed_eq_params_dict"
-                    " and observed_pinn_in_dict and observed_values_dict"
-                )
-            if observed_eq_params_dict is None:
-                observed_eq_params_dict = {
-                    k: None for k in observed_pinn_in_dict.keys()
-                }
-
-            keys = dict(
-                zip(
-                    observed_pinn_in_dict.keys(),
-                    jax.random.split(key, len(observed_pinn_in_dict)),
-                )
+        if self.observed_eq_params_dict is None:
+            self.observed_eq_params_dict = {
+                k: {} for k in self.observed_pinn_in_dict.keys()
+            }
+        elif self.observed_pinn_in_dict.keys() != self.observed_eq_params_dict.keys():
+            raise ValueError(
+                f"Keys must be the same in observed_eq_params_dict"
+                f" and observed_pinn_in_dict and observed_values_dict"
             )
-            self.data_gen_obs = jax.tree_util.tree_map(
-                lambda k, pinn_in, val, eq_params: (
-                    DataGeneratorObservations(
-                        k, obs_batch_size, pinn_in, val, eq_params
-                    )
-                    if pinn_in is not None
-                    else None
-                ),
-                keys,
-                observed_pinn_in_dict,
-                observed_values_dict,
-                observed_eq_params_dict,
+
+        keys = dict(
+            zip(
+                self.observed_pinn_in_dict.keys(),
+                jax.random.split(key, len(self.observed_pinn_in_dict)),
             )
+        )
+        self.data_gen_obs = jax.tree_util.tree_map(
+            lambda k, pinn_in, val, eq_params: (
+                DataGeneratorObservations(
+                    k, self.obs_batch_size, pinn_in, val, eq_params
+                )
+                if pinn_in is not None
+                else None
+            ),
+            keys,
+            self.observed_pinn_in_dict,
+            self.observed_values_dict,
+            self.observed_eq_params_dict,
+        )
 
-    def obs_batch(self):
+    def obs_batch(self) -> tuple["DataGeneratorObservationsMultiPINNs", PyTree]:
         """
         Returns a dictionary of DataGeneratorObservations.obs_batch with keys
         from `u_dict`
         """
-        return jax.tree_util.tree_map(
+        data_gen_and_batch_pytree = jax.tree_util.tree_map(
             lambda a: a.get_batch() if a is not None else {},
             self.data_gen_obs,
             is_leaf=lambda x: isinstance(x, DataGeneratorObservations),
         )  # note the is_leaf note to traverse the DataGeneratorObservations and
         # thus to be able to call the method on the element(s) of
         # self.data_gen_obs which are not None
+        new_attribute = jax.tree_util.tree_map(
+            lambda a: a[0],
+            data_gen_and_batch_pytree,
+            is_leaf=lambda x: isinstance(x, tuple),
+        )
+        new = eqx.tree_at(lambda m: m.data_gen_obs, self, new_attribute)
+        batches = jax.tree_util.tree_map(
+            lambda a: a[1],
+            data_gen_and_batch_pytree,
+            is_leaf=lambda x: isinstance(x, tuple),
+        )
 
-    def get_batch(self):
+        return new, batches
+
+    def get_batch(self) -> tuple["DataGeneratorObservationsMultiPINNs", PyTree]:
         """
         Generic method to return a batch
         """
         return self.obs_batch()
-
-    def tree_flatten(self):
-        # because a dict with "str" keys cannot go in the children (jittable)
-        # attributes, we need to separate it in two and recreate the zip in the
-        # tree_unflatten
-        children = self.data_gen_obs.values()
-        aux_data = {
-            "obs_batch_size": self.obs_batch_size,
-            "data_gen_obs_keys": self.data_gen_obs.keys(),
-        }
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        (data_gen_obs_values) = children
-        obj = cls(
-            observed_pinn_in_dict=None,
-            observed_values_dict=None,
-            data_gen_obs_exists=True,
-            obs_batch_size=aux_data["obs_batch_size"],
-        )
-        obj.data_gen_obs = dict(zip(aux_data["data_gen_obs_keys"], data_gen_obs_values))
-        return obj