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

jinns/loss/_loss_utils.py

@@ -6,42 +6,42 @@ from __future__ import (
     annotations,
 )  # https://docs.python.org/3/library/typing.html#constant
 
-from typing import TYPE_CHECKING, Callable, Dict
+from typing import TYPE_CHECKING, Callable, TypeGuard
+from types import EllipsisType
 import jax
 import jax.numpy as jnp
 from jax import vmap
-import equinox as eqx
-from jaxtyping import Float, Array, PyTree
+from jaxtyping import Float, Array
 
 from jinns.loss._boundary_conditions import (
     _compute_boundary_loss,
 )
 from jinns.utils._utils import _subtract_with_check, get_grid
-from jinns.data._DataGenerators import append_obs_batch, make_cartesian_product
+from jinns.data._utils import make_cartesian_product
 from jinns.parameters._params import _get_vmap_in_axes_params
 from jinns.nn._pinn import PINN
 from jinns.nn._spinn import SPINN
 from jinns.nn._hyperpinn import HyperPINN
-from jinns.data._Batchs import *
-from jinns.parameters._params import Params, ParamsDict
+from jinns.data._Batchs import PDEStatioBatch, PDENonStatioBatch
+from jinns.parameters._params import Params
 
 if TYPE_CHECKING:
-    from jinns.utils._types import *
+    from jinns.utils._types import BoundaryConditionFun
+    from jinns.nn._abstract_pinn import AbstractPINN
 
 
 def dynamic_loss_apply(
-    dyn_loss: DynamicLoss,
-    u: eqx.Module,
+    dyn_loss: Callable,
+    u: AbstractPINN,
     batch: (
-        Float[Array, "batch_size 1"]
-        | Float[Array, "batch_size dim"]
-        | Float[Array, "batch_size 1+dim"]
+        Float[Array, " batch_size 1"]
+        | Float[Array, " batch_size dim"]
+        | Float[Array, " batch_size 1+dim"]
     ),
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
-    loss_weight: float | Float[Array, "dyn_loss_dimension"],
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
     u_type: PINN | HyperPINN | None = None,
-) -> float:
+) -> Float[Array, " "]:
     """
     Sometimes when u is a lambda function a or dict we do not have access to
     its type here, hence the last argument
@@ -49,16 +49,18 @@ def dynamic_loss_apply(
     if u_type == PINN or u_type == HyperPINN or isinstance(u, (PINN, HyperPINN)):
         v_dyn_loss = vmap(
             lambda batch, params: dyn_loss(
-                batch, u, params  # we must place the params at the end
+                batch,
+                u,
+                params,  # we must place the params at the end
             ),
             vmap_axes,
             0,
         )
         residuals = v_dyn_loss(batch, params)
-        mse_dyn_loss = jnp.mean(jnp.sum(loss_weight * residuals**2, axis=-1))
+        mse_dyn_loss = jnp.mean(jnp.sum(residuals**2, axis=-1))
     elif u_type == SPINN or isinstance(u, SPINN):
         residuals = dyn_loss(batch, u, params)
-        mse_dyn_loss = jnp.mean(jnp.sum(loss_weight * residuals**2, axis=-1))
+        mse_dyn_loss = jnp.mean(jnp.sum(residuals**2, axis=-1))
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
 
@@ -66,18 +68,17 @@ def dynamic_loss_apply(
 
 
 def normalization_loss_apply(
-    u: eqx.Module,
+    u: AbstractPINN,
     batches: (
-        tuple[Float[Array, "nb_norm_samples dim"]]
+        tuple[Float[Array, " nb_norm_samples dim"]]
         | tuple[
-            Float[Array, "nb_norm_time_slices 1"], Float[Array, "nb_norm_samples dim"]
+            Float[Array, " nb_norm_time_slices 1"], Float[Array, " nb_norm_samples dim"]
         ]
     ),
-    params: Params | ParamsDict,
-    vmap_axes_params: tuple[int | None, ...],
-    norm_weights: Float[Array, "nb_norm_samples"],
-    loss_weight: float,
-) -> float:
+    params: Params[Array],
+    vmap_axes_params: tuple[Params[int | None] | None],
+    norm_weights: Float[Array, " nb_norm_samples"],
+) -> Float[Array, " "]:
     """
     Note the squeezing on each result. We expect unidimensional *PINN since
     they represent probability distributions
@@ -92,12 +93,10 @@ def normalization_loss_apply(
             res = v_u(*batches, params)
             assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
             # Monte-Carlo integration using importance sampling
-            mse_norm_loss = loss_weight * (
-                jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
-            )
+            mse_norm_loss = jnp.abs(jnp.mean(res.squeeze() * norm_weights) - 1) ** 2
         else:
             # NOTE this cartesian product is costly
-            batches = make_cartesian_product(
+            batch_cart_prod = make_cartesian_product(
                 batches[0],
                 batches[1],
             ).reshape(batches[0].shape[0], batches[1].shape[0], -1)
@@ -108,11 +107,11 @@ def normalization_loss_apply(
                 ),
                 in_axes=(0,) + vmap_axes_params,
             )
-            res = v_u(batches, params)
+            res = v_u(batch_cart_prod, params)
             assert res.shape[-1] == 1, "norm loss expects unidimensional *PINN"
             # For all times t, we perform an integration. Then we average the
             # losses over times.
-            mse_norm_loss = loss_weight * jnp.mean(
+            mse_norm_loss = jnp.mean(
                 jnp.abs(jnp.mean(res.squeeze() * norm_weights, axis=-1) - 1) ** 2
             )
     elif isinstance(u, SPINN):
@@ -120,8 +119,7 @@ def normalization_loss_apply(
             res = u(*batches, params)
             assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             mse_norm_loss = (
-                loss_weight
-                * jnp.abs(
+                jnp.abs(
                     jnp.mean(
                         res.squeeze(),
                     )
@@ -141,11 +139,11 @@ def normalization_loss_apply(
             )
             assert res.shape[-1] == 1, "norm loss expects unidimensional *SPINN"
             # the outer mean() below is for the times stamps
-            mse_norm_loss = loss_weight * jnp.mean(
+            mse_norm_loss = jnp.mean(
                 jnp.abs(
                     jnp.mean(
                         res.squeeze(),
-                        axis=(d + 1 for d in range(res.ndim - 2)),
+                        axis=list(d + 1 for d in range(res.ndim - 2)),
                     )
                     * norm_weights
                     - 1
@@ -159,18 +157,33 @@ def normalization_loss_apply(
 
 
 def boundary_condition_apply(
-    u: eqx.Module,
+    u: AbstractPINN,
     batch: PDEStatioBatch | PDENonStatioBatch,
-    params: Params | ParamsDict,
-    omega_boundary_fun: Callable,
-    omega_boundary_condition: str,
-    omega_boundary_dim: int,
-    loss_weight: float | Float[Array, "boundary_cond_dim"],
-) -> float:
-
+    params: Params[Array],
+    omega_boundary_fun: BoundaryConditionFun | dict[str, BoundaryConditionFun],
+    omega_boundary_condition: str | dict[str, str],
+    omega_boundary_dim: slice | dict[str, slice],
+) -> Float[Array, " "]:
+    assert batch.border_batch is not None
     vmap_in_axes = (0,) + _get_vmap_in_axes_params(batch.param_batch_dict, params)
 
-    if isinstance(omega_boundary_fun, dict):
+    def _check_tuple_of_dict(
+        val,
+    ) -> TypeGuard[
+        tuple[
+            dict[str, BoundaryConditionFun],
+            dict[str, BoundaryConditionFun],
+            dict[str, BoundaryConditionFun],
+        ]
+    ]:
+        return all(isinstance(x, dict) for x in val)
+
+    omega_boundary_dicts = (
+        omega_boundary_condition,
+        omega_boundary_fun,
+        omega_boundary_dim,
+    )
+    if _check_tuple_of_dict(omega_boundary_dicts):
         # We must create the facet tree dictionary as we do not have the
         # enumerate from the for loop to pass the id integer
         if batch.border_batch.shape[-1] == 2:
@@ -186,16 +199,13 @@ def boundary_condition_apply(
                 None
                 if c is None
                 else jnp.mean(
-                    loss_weight
-                    * _compute_boundary_loss(
-                        c, f, batch, u, params, fa, d, vmap_in_axes
-                    )
+                    _compute_boundary_loss(c, f, batch, u, params, fa, d, vmap_in_axes)
                 )
             ),
-            omega_boundary_condition,
-            omega_boundary_fun,
+            omega_boundary_dicts[0],  # omega_boundary_condition,
+            omega_boundary_dicts[1],  # omega_boundary_fun,
             facet_tree,
-            omega_boundary_dim,
+            omega_boundary_dicts[2],  # omega_boundary_dim,
             is_leaf=lambda x: x is None,
         )  # when exploring leaves with None value (no condition) the returned
         # mse is None and we get rid of the None leaves of b_losses_by_facet
@@ -206,15 +216,14 @@ def boundary_condition_apply(
         facet_tuple = tuple(f for f in range(batch.border_batch.shape[-1]))
         b_losses_by_facet = jax.tree_util.tree_map(
             lambda fa: jnp.mean(
-                loss_weight
-                * _compute_boundary_loss(
-                    omega_boundary_condition,
-                    omega_boundary_fun,
+                _compute_boundary_loss(
+                    omega_boundary_dicts[0],  # type: ignore -> need TypeIs from 3.13
+                    omega_boundary_dicts[1],  # type: ignore -> need TypeIs from 3.13
                     batch,
                     u,
                     params,
                     fa,
-                    omega_boundary_dim,
+                    omega_boundary_dicts[2],  # type: ignore -> need TypeIs from 3.13
                     vmap_in_axes,
                 )
             ),
@@ -227,26 +236,23 @@ def boundary_condition_apply(
 
 
 def observations_loss_apply(
-    u: eqx.Module,
-    batches: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
-    observed_values: Float[Array, "batch_size observation_dim"],
-    loss_weight: float | Float[Array, "observation_dim"],
-    obs_slice: slice,
-) -> float:
-    # TODO implement for SPINN
+    u: AbstractPINN,
+    batch: Float[Array, " obs_batch_size input_dim"],
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
+    observed_values: Float[Array, " obs_batch_size observation_dim"],
+    obs_slice: EllipsisType | slice | None,
+) -> Float[Array, " "]:
     if isinstance(u, (PINN, HyperPINN)):
         v_u = vmap(
             lambda *args: u(*args)[u.slice_solution],
             vmap_axes,
             0,
         )
-        val = v_u(*batches, params)[:, obs_slice]
+        val = v_u(batch, params)[:, obs_slice]
         mse_observation_loss = jnp.mean(
             jnp.sum(
-                loss_weight
-                * _subtract_with_check(
+                _subtract_with_check(
                     observed_values, val, cause="user defined observed_values"
                 )
                 ** 2,
@@ -261,15 +267,15 @@ def observations_loss_apply(
 
 
 def initial_condition_apply(
-    u: eqx.Module,
-    omega_batch: Float[Array, "dimension"],
-    params: Params | ParamsDict,
-    vmap_axes: tuple[int | None, ...],
+    u: AbstractPINN,
+    omega_batch: Float[Array, " dimension"],
+    params: Params[Array],
+    vmap_axes: tuple[int, Params[int | None] | None],
     initial_condition_fun: Callable,
-    loss_weight: float | Float[Array, "initial_condition_dimension"],
-) -> float:
+    t0: Float[Array, " 1"],
+) -> Float[Array, " "]:
     n = omega_batch.shape[0]
-    t0_omega_batch = jnp.concatenate([jnp.zeros((n, 1)), omega_batch], axis=1)
+    t0_omega_batch = jnp.concatenate([t0 * jnp.ones((n, 1)), omega_batch], axis=1)
     if isinstance(u, (PINN, HyperPINN)):
         v_u_t0 = vmap(
             lambda t0_x, params: _subtract_with_check(
@@ -285,7 +291,7 @@ def initial_condition_apply(
         # dimension as params to be able to vmap.
         # Recall that by convention:
         # param_batch_dict = times_batch_size * omega_batch_size
-        mse_initial_condition = jnp.mean(jnp.sum(loss_weight * res**2, axis=-1))
+        mse_initial_condition = jnp.mean(jnp.sum(res**2, axis=-1))
     elif isinstance(u, SPINN):
         values = lambda t_x: u(
             t_x,
@@ -298,107 +304,7 @@ def initial_condition_apply(
             v_ini,
             cause="Output of initial_condition_fun",
         )
-        mse_initial_condition = jnp.mean(jnp.sum(loss_weight * res**2, axis=-1))
+        mse_initial_condition = jnp.mean(jnp.sum(res**2, axis=-1))
     else:
         raise ValueError(f"Bad type for u. Got {type(u)}, expected PINN or SPINN")
     return mse_initial_condition
-
-
-def constraints_system_loss_apply(
-    u_constraints_dict: Dict,
-    batch: ODEBatch | PDEStatioBatch | PDENonStatioBatch,
-    params_dict: ParamsDict,
-    loss_weights: Dict[str, float | Array],
-    loss_weight_struct: PyTree,
-):
-    """
-    Same function for systemlossODE and systemlossPDE!
-    """
-    # Transpose so we have each u_dict as outer structure and the
-    # associated loss_weight as inner structure
-    loss_weights_T = jax.tree_util.tree_transpose(
-        jax.tree_util.tree_structure(loss_weight_struct),
-        jax.tree_util.tree_structure(loss_weights["initial_condition"]),
-        loss_weights,
-    )
-
-    if isinstance(params_dict.nn_params, dict):
-
-        def apply_u_constraint(
-            u_constraint, nn_params, eq_params, loss_weights_for_u, obs_batch_u
-        ):
-            res_dict_for_u = u_constraint.evaluate(
-                Params(
-                    nn_params=nn_params,
-                    eq_params=eq_params,
-                ),
-                append_obs_batch(batch, obs_batch_u),
-            )[1]
-            res_dict_ponderated = jax.tree_util.tree_map(
-                lambda w, l: w * l, res_dict_for_u, loss_weights_for_u
-            )
-            return res_dict_ponderated
-
-        # Note in the case of multiple PINNs, batch.obs_batch_dict is a dict
-        # with keys corresponding to the PINN and value correspondinf to an
-        # original obs_batch_dict. Hence the tree mapping also interates over
-        # batch.obs_batch_dict
-        res_dict = jax.tree_util.tree_map(
-            apply_u_constraint,
-            u_constraints_dict,
-            params_dict.nn_params,
-            (
-                params_dict.eq_params
-                if params_dict.eq_params.keys() == params_dict.nn_params.keys()
-                else {k: params_dict.eq_params for k in params_dict.nn_params.keys()}
-            ),  # this manipulation is needed since we authorize eq_params not to have the same structure as nn_params in ParamsDict
-            loss_weights_T,
-            batch.obs_batch_dict,
-            is_leaf=lambda x: (
-                not isinstance(x, dict)  # to not traverse more than the first
-                # outer dict of the pytrees passed to the function. This will
-                # work because u_constraints_dict is a dict of Losses, and it
-                # thus stops the traversing of other dict too
-            ),
-        )
-    # TODO try to get rid of this condition?
-    else:
-
-        def apply_u_constraint(u_constraint, loss_weights_for_u, obs_batch_u):
-            res_dict_for_u = u_constraint.evaluate(
-                params_dict,
-                append_obs_batch(batch, obs_batch_u),
-            )[1]
-            res_dict_ponderated = jax.tree_util.tree_map(
-                lambda w, l: w * l, res_dict_for_u, loss_weights_for_u
-            )
-            return res_dict_ponderated
-
-        res_dict = jax.tree_util.tree_map(
-            apply_u_constraint, u_constraints_dict, loss_weights_T, batch.obs_batch_dict
-        )
-
-    # Transpose back so we have mses as outer structures and their values
-    # for each u_dict as inner structures. The tree_leaves transforms the
-    # inner structure into a list so we can catch is as leaf it the
-    # tree_map below
-    res_dict = jax.tree_util.tree_transpose(
-        jax.tree_util.tree_structure(
-            jax.tree_util.tree_leaves(loss_weights["initial_condition"])
-        ),
-        jax.tree_util.tree_structure(loss_weight_struct),
-        res_dict,
-    )
-    # For each mse, sum their values on each u_dict
-    res_dict = jax.tree_util.tree_map(
-        lambda mse: jax.tree_util.tree_reduce(
-            lambda x, y: x + y, jax.tree_util.tree_leaves(mse)
-        ),
-        res_dict,
-        is_leaf=lambda x: isinstance(x, list),
-    )
-    # Total loss
-    total_loss = jax.tree_util.tree_reduce(
-        lambda x, y: x + y, jax.tree_util.tree_leaves(res_dict)
-    )
-    return total_loss, res_dict

jinns/loss/_loss_weight_updates.py

@@ -0,0 +1,202 @@
+"""
+A collection of specific weight update schemes in jinns
+"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from jaxtyping import Array, Key
+import jax.numpy as jnp
+import jax
+import equinox as eqx
+
+if TYPE_CHECKING:
+    from jinns.loss._loss_weights import AbstractLossWeights
+    from jinns.utils._types import AnyLossComponents
+
+
+def soft_adapt(
+    loss_weights: AbstractLossWeights,
+    iteration_nb: int,
+    loss_terms: AnyLossComponents,
+    stored_loss_terms: AnyLossComponents,
+) -> Array:
+    r"""
+    Implement the simple strategy given in
+    https://docs.nvidia.com/deeplearning/physicsnemo/physicsnemo-sym/user_guide/theory/advanced_schemes.html#softadapt
+
+    $$
+    w_j(i)= \frac{\exp(\frac{L_j(i)}{L_j(i-1)+\epsilon}-\mu(i))}
+    {\sum_{k=1}^{n_{loss}}\exp(\frac{L_k(i)}{L_k(i-1)+\epsilon}-\mu(i)}
+    $$
+
+    Note that since None is not treated as a leaf by jax tree.util functions,
+    we naturally avoid None components from loss_terms, stored_loss_terms etc.!
+    """
+
+    def do_nothing(loss_weights, _, __):
+        return jnp.array(
+            jax.tree.leaves(loss_weights, is_leaf=eqx.is_inexact_array), dtype=float
+        )
+
+    def soft_adapt_(_, loss_terms, stored_loss_terms):
+        ratio_pytree = jax.tree.map(
+            lambda lt, slt: lt / (slt[iteration_nb - 1] + 1e-6),
+            loss_terms,
+            stored_loss_terms,
+        )
+        mu = jax.tree.reduce(jnp.maximum, ratio_pytree, initializer=jnp.array(-jnp.inf))
+        ratio_pytree = jax.tree.map(lambda r: r - mu, ratio_pytree)
+        ratio_leaves = jax.tree.leaves(ratio_pytree)
+        return jax.nn.softmax(jnp.array(ratio_leaves))
+
+    return jax.lax.cond(
+        iteration_nb == 0,
+        lambda op: do_nothing(*op),
+        lambda op: soft_adapt_(*op),
+        (loss_weights, loss_terms, stored_loss_terms),
+    )
+
+
+def ReLoBRaLo(
+    loss_weights: AbstractLossWeights,
+    iteration_nb: int,
+    loss_terms: AnyLossComponents,
+    stored_loss_terms: AnyLossComponents,
+    key: Key,
+    decay_factor: float = 0.9,
+    tau: float = 1,  ## referred to as temperature in the article
+    p: float = 0.9,
+):
+    r"""
+    Implementing the extension of softadapt: Relative Loss Balancing with random LookBack
+    """
+    n_loss = len(jax.tree.leaves(loss_terms))  # number of loss terms
+    epsilon = 1e-6
+
+    def do_nothing(loss_weights, _):
+        return jnp.array(
+            jax.tree.leaves(loss_weights, is_leaf=eqx.is_inexact_array), dtype=float
+        )
+
+    def compute_softmax_weights(current, reference):
+        ratio_pytree = jax.tree.map(
+            lambda lt, ref: lt / (ref + epsilon),
+            current,
+            reference,
+        )
+        mu = jax.tree.reduce(jnp.maximum, ratio_pytree, initializer=-jnp.inf)
+        ratio_pytree = jax.tree.map(lambda r: r - mu, ratio_pytree)
+        ratio_leaves = jax.tree.leaves(ratio_pytree)
+        return jax.nn.softmax(jnp.array(ratio_leaves))
+
+    def soft_adapt_prev(stored_loss_terms):
+        # ω_j(i-1)
+        prev_terms = jax.tree.map(lambda slt: slt[iteration_nb - 1], stored_loss_terms)
+        prev_prev_terms = jax.tree.map(
+            lambda slt: slt[iteration_nb - 2], stored_loss_terms
+        )
+        return compute_softmax_weights(prev_terms, prev_prev_terms)
+
+    def look_back(loss_terms, stored_loss_terms):
+        # ω̂_j^(i,0)
+        initial_terms = jax.tree.map(lambda slt: tau * slt[0], stored_loss_terms)
+        weights = compute_softmax_weights(loss_terms, initial_terms)
+        return n_loss * weights
+
+    def soft_adapt_current(loss_terms, stored_loss_terms):
+        # ω_j(i)
+        prev_terms = jax.tree.map(lambda slt: slt[iteration_nb - 1], stored_loss_terms)
+        return compute_softmax_weights(loss_terms, prev_terms)
+
+    # Bernoulli variable for random lookback
+    rho = jax.random.bernoulli(key, p).astype(float)
+
+    # Base case for first iteration
+    def first_iter_case(_):
+        return do_nothing(loss_weights, None)
+
+    # Case for iteration >= 1
+    def subsequent_iter_case(_):
+        # Compute historical weights
+        def hist_weights_case1(_):
+            return soft_adapt_current(loss_terms, stored_loss_terms)
+
+        def hist_weights_case2(_):
+            return rho * soft_adapt_prev(stored_loss_terms) + (1 - rho) * look_back(
+                loss_terms, stored_loss_terms
+            )
+
+        loss_weights_hist = jax.lax.cond(
+            iteration_nb < 2,
+            hist_weights_case1,
+            hist_weights_case2,
+            None,
+        )
+
+        # Compute and return final weights
+        adaptive_weights = soft_adapt_current(loss_terms, stored_loss_terms)
+        return decay_factor * loss_weights_hist + (1 - decay_factor) * adaptive_weights
+
+    return jax.lax.cond(
+        iteration_nb == 0,
+        first_iter_case,
+        subsequent_iter_case,
+        None,
+    )
+
+
+def lr_annealing(
+    loss_weights: AbstractLossWeights,
+    grad_terms: AnyLossComponents,
+    decay_factor: float = 0.9,  # 0.9 is the recommended value from the article
+    eps: float = 1e-6,
+) -> Array:
+    r"""
+    Implementation of the Learning rate annealing
+    Algorithm 1 in the paper UNDERSTANDING AND MITIGATING GRADIENT PATHOLOGIES IN PHYSICS-INFORMED NEURAL NETWORKS
+
+    (a) Compute $\hat{\lambda}_i$ by
+    $$
+        \hat{\lambda}_i = \frac{\max_{\theta}\{|\nabla_\theta \mathcal{L}_r (\theta_n)|\}}{mean(|\nabla_\theta \mathcal{L}_i (\theta_n)|)}, \quad i=1,\dots, M,
+    $$
+
+    (b) Update the weights $\lambda_i$ using a moving average of the form
+    $$
+        \lambda_i = (1-\alpha) \lambda_{i-1} + \alpha \hat{\lambda}_i, \quad i=1, \dots, M.
+    $$
+
+    Note that since None is not treated as a leaf by jax tree.util functions,
+    we naturally avoid None components from loss_terms, stored_loss_terms etc.!
+
+    """
+    assert hasattr(grad_terms, "dyn_loss")
+    dyn_loss_grads = getattr(grad_terms, "dyn_loss")
+    data_fit_grads = [
+        getattr(grad_terms, att) if hasattr(grad_terms, att) else None
+        for att in ["norm_loss", "boundary_loss", "observations", "initial_condition"]
+    ]
+
+    dyn_loss_grads_leaves = jax.tree.leaves(
+        dyn_loss_grads,
+        is_leaf=eqx.is_inexact_array,
+    )
+
+    max_dyn_loss_grads = jnp.max(
+        jnp.stack([jnp.max(jnp.abs(g)) for g in dyn_loss_grads_leaves])
+    )
+
+    mean_gradients = [
+        jnp.mean(jnp.stack([jnp.abs(jnp.mean(g)) for g in jax.tree.leaves(t)]))
+        for t in data_fit_grads
+        if t is not None and jax.tree.leaves(t)
+    ]
+
+    lambda_hat = max_dyn_loss_grads / (jnp.array(mean_gradients) + eps)
+    old_weights = jnp.array(
+        jax.tree.leaves(
+            loss_weights,
+        )
+    )
+
+    new_weights = (1 - decay_factor) * old_weights[1:] + decay_factor * lambda_hat
+    return jnp.hstack([old_weights[0], new_weights])