diffbio 0.1.4__tar.gz → 0.1.6__tar.gz

{diffbio-0.1.4 → diffbio-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.4
+Version: 0.1.6
 Summary: End-to-end differentiable bioinformatics for JAX/Flax NNX: alignment, mapping, assembly, variant calling, RNA-seq, single-cell, epigenomics, CRISPR, metabolomics, multi-omics, protein and RNA structure, molecular dynamics and drug-discovery operators composed into trainable pipelines on datarax, artifex, opifex and calibrax
 Project-URL: Bug Tracker, https://github.com/avitai/DiffBio/issues
 Project-URL: Documentation, https://diffbio.readthedocs.io
@@ -45,22 +45,22 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.14,>=3.12
 Requires-Dist: anndata>=0.9.1
-Requires-Dist: avitai-artifex>=0.1.5
+Requires-Dist: avitai-artifex>=0.1.9
 Requires-Dist: beartype>=0.14.1
 Requires-Dist: biopython>=1.81
-Requires-Dist: calibrax>=0.1.5
+Requires-Dist: calibrax>=0.1.6
 Requires-Dist: chex>=0.1.7
-Requires-Dist: datarax>=0.1.7
+Requires-Dist: datarax>=0.1.11
 Requires-Dist: flax>=0.12.9
 Requires-Dist: h5py>=3.7
 Requires-Dist: jax-md>=0.2.27
 Requires-Dist: jax>=0.11.1
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: numpy>=1.24
-Requires-Dist: opifex>=0.2.2
+Requires-Dist: opifex>=0.2.7
 Requires-Dist: optax>=0.2.8
 Requires-Dist: scipy>=1.10
-Requires-Dist: substrax>=0.1.5
+Requires-Dist: substrax>=0.1.9
 Provides-Extra: all
 Requires-Dist: bandit[toml]>=1.8.6; extra == 'all'
 Requires-Dist: beartype>=0.14.1; extra == 'all'
@@ -353,6 +353,7 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 
 ```python
 from diffbio.utils import (
+from substrax.optim import OptimizerConfig
     Trainer,
     TrainingConfig,
     cross_entropy_loss,
@@ -371,10 +372,11 @@ inputs, targets = create_synthetic_training_data(
 
 # Configure training
 config = TrainingConfig(
-    learning_rate=1e-3,
+    optimizer=OptimizerConfig(
+        optimizer_type="adam", learning_rate=1e-3, gradient_clip_norm=1.0
+    ),
     num_epochs=50,
     log_every=10,
-    grad_clip_norm=1.0,
 )
 
 # Create trainer

{diffbio-0.1.4 → diffbio-0.1.6}/README.md

@@ -139,6 +139,7 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 
 ```python
 from diffbio.utils import (
+from substrax.optim import OptimizerConfig
     Trainer,
     TrainingConfig,
     cross_entropy_loss,
@@ -157,10 +158,11 @@ inputs, targets = create_synthetic_training_data(
 
 # Configure training
 config = TrainingConfig(
-    learning_rate=1e-3,
+    optimizer=OptimizerConfig(
+        optimizer_type="adam", learning_rate=1e-3, gradient_clip_norm=1.0
+    ),
     num_epochs=50,
     log_every=10,
-    grad_clip_norm=1.0,
 )
 
 # Create trainer

{diffbio-0.1.4 → diffbio-0.1.6}/pyproject.toml

@@ -37,15 +37,16 @@ classifiers = [
 ]
 dependencies = [
   # Datarax - operator, source, and pipeline contracts
-  "datarax>=0.1.7",
+  "datarax>=0.1.11",
   # Artifex - modeling, modality, and transformer substrate
-  "avitai-artifex>=0.1.5",
+  "avitai-artifex>=0.1.9",
   # Opifex - scientific ML, operator learning, and advanced optimization
-  "opifex>=0.2.2",
+  "opifex>=0.2.7",
   # Calibrax - benchmarking, comparison, profiling, and regression control
-  "calibrax>=0.1.5",
-  # Substrax - device identity, meshes and checkpoint store shared by the siblings
-  "substrax>=0.1.5",
+  "calibrax>=0.1.6",
+  # Substrax - device identity, meshes, runtime settings and checkpoint store shared by the
+  # siblings; 0.1.6 adds substrax.runtime, substrax.artifacts and substrax.testing.
+  "substrax>=0.1.9",
   # JAX ecosystem
   "beartype>=0.14.1",
   "chex>=0.1.7",
@@ -70,7 +71,7 @@ license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.4"
+version = "0.1.6"
 
 [project.optional-dependencies]
 all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/core/base_operators.py

@@ -25,11 +25,13 @@ from datarax.core.config import OperatorConfig
 from datarax.core.operator import OperatorModule
 from flax import nnx
 from jaxtyping import Array, Float, Int, PyTree
+from substrax.rng import key_from
 
 from diffbio.constants import DEFAULT_TEMPERATURE
 from diffbio.core.graph_utils import scatter_aggregate as scatter_aggregate_fn
 from diffbio.core.soft_ops import sorting as soft_sorting
-from diffbio.utils.nn_utils import ensure_rngs, get_rng_key, init_learnable_param
+from diffbio.utils.nn_utils import init_learnable_param
+
 
 __all__ = [
     "TemperatureOperator",
@@ -139,7 +141,7 @@ class TemperatureOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Base apply method - should be overridden by subclasses."""
@@ -233,7 +235,7 @@ class SequenceOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Base apply method - should be overridden by subclasses."""
@@ -258,7 +260,7 @@ class EncoderDecoderOperator(OperatorModule):
         self,
         config: OperatorConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize EncoderDecoderOperator.
@@ -272,12 +274,13 @@ class EncoderDecoderOperator(OperatorModule):
 
         self.latent_dim = getattr(config, "latent_dim", 10)
         self.hidden_dim = getattr(config, "hidden_dim", 64)
-        self.rngs = ensure_rngs(rngs)
+        self.rngs = rngs
 
     def reparameterize(
         self,
         mean: Float[Array, "... latent_dim"],
         log_var: Float[Array, "... latent_dim"],
+        key: jax.Array | None = None,
     ) -> Float[Array, "... latent_dim"]:
         """Sample from latent distribution using reparameterization trick.
 
@@ -288,11 +291,18 @@ class EncoderDecoderOperator(OperatorModule):
         Args:
             mean: Mean of the latent distribution.
             log_var: Log variance of the latent distribution.
+            key: The key to draw epsilon from. ``apply`` passes the record's key, so a
+                record's sample follows the record; a training loss such as
+                ``compute_elbo_loss`` passes none and draws from the operator's ``sample``
+                stream, which advances on every call.
 
         Returns:
             Sampled latent representation.
         """
-        key = get_rng_key(self.rngs, "sample", fallback_seed=0)
+        if key is None:
+            key = key_from(
+                self.rngs, streams=("sample", "default"), context=f"{type(self).__name__} sampling"
+            )
         std = jnp.exp(0.5 * log_var)
         epsilon = jax.random.normal(key, mean.shape)
         return mean + std * epsilon
@@ -345,7 +355,7 @@ class EncoderDecoderOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Base apply method - should be overridden by subclasses."""
@@ -442,7 +452,7 @@ class GraphOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Base apply method - should be overridden by subclasses."""
@@ -467,7 +477,7 @@ class HMMOperator(OperatorModule):
         self,
         config: OperatorConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize HMMOperator.
@@ -484,20 +494,20 @@ class HMMOperator(OperatorModule):
         self.temperature = getattr(config, "temperature", DEFAULT_TEMPERATURE)
 
         # Initialize HMM parameters
-        rngs = ensure_rngs(rngs)
 
+        context = f"{type(self).__name__} parameters"
         # Transition logits (will be normalized via log_softmax)
-        key = get_rng_key(rngs, "params", fallback_seed=0)
+        key = key_from(rngs, streams=("params", "default"), context=context)
         init_trans = jax.random.normal(key, (self.num_states, self.num_states)) * 0.1
         self.log_transition_params = nnx.Param(init_trans)
 
         # Emission logits
-        key = get_rng_key(rngs, "params", fallback_seed=1)
+        key = key_from(rngs, streams=("params", "default"), context=context)
         init_emit = jax.random.normal(key, (self.num_states, self.num_emissions)) * 0.1
         self.log_emission_params = nnx.Param(init_emit)
 
         # Initial state logits
-        key = get_rng_key(rngs, "params", fallback_seed=2)
+        key = key_from(rngs, streams=("params", "default"), context=context)
         init_initial = jax.random.normal(key, (self.num_states,)) * 0.1
         self.log_initial_params = nnx.Param(init_initial)
 
@@ -593,7 +603,7 @@ class HMMOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Base apply method - should be overridden by subclasses."""

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/core/data_types.py

@@ -6,6 +6,7 @@ the expected interfaces and data structures across the DiffBio codebase.
 
 from typing import Any, Protocol, TypedDict, runtime_checkable
 
+import jax
 from jaxtyping import Array, Float
 
 # =============================================================================
@@ -166,7 +167,7 @@ class DifferentiableOperator(Protocol):
         data: dict[str, Any],
         state: StateDict,
         metadata: MetadataDict,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> OperatorOutput:
         """Apply the operator to input data.
@@ -175,7 +176,7 @@ class DifferentiableOperator(Protocol):
             data: Input data dictionary.
             state: Element state.
             metadata: Element metadata.
-            random_params: Random parameters for stochastic operations.
+            key: The record's PRNG key; a stochastic operator draws from it.
             stats: Statistics dictionary.
 
         Returns:

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/core/neural_components.py

@@ -27,17 +27,10 @@ from typing import Literal
 import jax
 import jax.numpy as jnp
 from artifex.generative_models.core.base import MLP
-from flax import nnx
-from jaxtyping import Array, Float, Int
-
-from diffbio.constants import DEFAULT_TEMPERATURE
-from diffbio.core.graph_utils import scatter_aggregate
-from diffbio.utils.nn_utils import get_rng_key
 
 # =============================================================================
 # Re-export from artifex (import when available, provide stubs otherwise)
 # =============================================================================
-
 from artifex.generative_models.core.layers.positional import (
     PositionalEncoding,
     RotaryPositionalEncoding as RoPE,
@@ -47,6 +40,12 @@ from artifex.generative_models.core.layers.residual import (
     Conv1DResidualBlock as ResidualBlock1D,
     Conv2DResidualBlock as ResidualBlock2D,
 )
+from flax import nnx
+from jaxtyping import Array, Float, Int
+from substrax.rng import key_from
+
+from diffbio.constants import DEFAULT_TEMPERATURE
+from diffbio.core.graph_utils import scatter_aggregate
 
 
 __all__ = [
@@ -115,7 +114,9 @@ class GumbelSoftmaxModule(nnx.Module):
         Returns:
             Samples of same shape as logits.
         """
-        key = get_rng_key(self.rngs, "dropout", fallback_seed=0)
+        key = key_from(
+            self.rngs, streams=("dropout", "default"), context="GumbelSoftmaxModule sampling"
+        )
         gumbel_noise = jax.random.gumbel(key, logits.shape)
         perturbed = (logits + gumbel_noise) / self.temperature
         soft_sample = jax.nn.softmax(perturbed, axis=-1)

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/core/uncertainty.py

@@ -161,7 +161,7 @@ class EnsembleUQOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,  # noqa: ARG002
+        key: jax.Array | None = None,  # noqa: ARG002
         stats: dict[str, Any] | None = None,  # noqa: ARG002
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Run ensemble forward passes and aggregate with uncertainty.
@@ -170,7 +170,7 @@ class EnsembleUQOperator(OperatorModule):
             data: Input data dict for the base operator.
             state: Element state (passed through).
             metadata: Element metadata (passed through).
-            random_params: Unused.
+            key: Unused.
             stats: Unused.
 
         Returns:
@@ -241,7 +241,7 @@ class MCSamplingUQOperator(OperatorModule):
         data: PyTree,
         state: PyTree,
         metadata: dict[str, Any] | None,
-        random_params: Any = None,  # noqa: ARG002
+        key: jax.Array | None = None,  # noqa: ARG002
         stats: dict[str, Any] | None = None,  # noqa: ARG002
     ) -> tuple[PyTree, PyTree, dict[str, Any] | None]:
         """Run Monte Carlo sampling to estimate empirical intervals.
@@ -250,7 +250,7 @@ class MCSamplingUQOperator(OperatorModule):
             data: Input data dict for the base operator.
             state: Element state (passed through).
             metadata: Element metadata (passed through).
-            random_params: Unused.
+            key: Unused.
             stats: Unused.
 
         Returns:

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/evaluation/adapters.py

@@ -135,7 +135,7 @@ class TaskAdapter:
         else:
             adapted = data_dict
 
-        result, _, _ = operator.apply(adapted, {}, None)
+        result, _, _ = operator.apply(adapted, {}, None, jax.random.key(self._seed))
         # Propagate retention info
         if "retention_weights" not in result and "quality_scores" in adapted:
             quality = adapted["quality_scores"]
@@ -160,7 +160,7 @@ class TaskAdapter:
             temperature=temperature,
         )
         operator = SoftKMeansClustering(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = operator.apply(data_dict, {}, None)
+        result, _, _ = operator.apply(data_dict, {}, None, jax.random.key(self._seed))
         return result
 
     def _run_de(self, data_dict: dict[str, Any], config: dict[str, Any]) -> dict[str, Any]:
@@ -178,7 +178,7 @@ class TaskAdapter:
             n_conditions=n_conditions,
         )
         pipeline = DifferentialExpressionPipeline(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = pipeline.apply(data_dict, {}, None)
+        result, _, _ = pipeline.apply(data_dict, {}, None, jax.random.key(self._seed))
 
         # Add gene names for grader extraction
         if "gene_names" not in result:
@@ -202,7 +202,7 @@ class TaskAdapter:
             n_clusters=n_clusters,
         )
         operator = DifferentiableHarmony(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = operator.apply(data_dict, {}, None)
+        result, _, _ = operator.apply(data_dict, {}, None, jax.random.key(self._seed))
         return result
 
     def _run_normalization(
@@ -222,7 +222,7 @@ class TaskAdapter:
             latent_dim=latent_dim,
         )
         operator = VAENormalizer(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = operator.apply(data_dict, {}, None)
+        result, _, _ = operator.apply(data_dict, {}, None, jax.random.key(self._seed))
         return result
 
     def _run_trajectory(self, data_dict: dict[str, Any], config: dict[str, Any]) -> dict[str, Any]:
@@ -240,7 +240,7 @@ class TaskAdapter:
             n_diffusion_components=n_diffusion_components,
         )
         operator = DifferentiablePseudotime(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = operator.apply(data_dict, {}, None)
+        result, _, _ = operator.apply(data_dict, {}, None, jax.random.key(self._seed))
         return result
 
     def _run_spatial_analysis(
@@ -260,7 +260,7 @@ class TaskAdapter:
             n_domains=n_domains,
         )
         operator = DifferentiableSpatialDomain(op_config, rngs=nnx.Rngs(self._seed))
-        result, _, _ = operator.apply(data_dict, {}, None)
+        result, _, _ = operator.apply(data_dict, {}, None, jax.random.key(self._seed))
         return result
 
     def _run_cell_annotation(

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/losses/statistical_losses.py

@@ -261,7 +261,7 @@ class HMMLikelihoodLoss(nnx.Module):
         n_states: int,
         n_emissions: int,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
     ):
         """Initialize the HMM loss.
 
@@ -272,9 +272,6 @@ class HMMLikelihoodLoss(nnx.Module):
         """
         super().__init__()
 
-        if rngs is None:
-            rngs = nnx.Rngs(0)
-
         self.n_states = n_states
         self.n_emissions = n_emissions
 

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/operators/_count_vae.py

@@ -4,14 +4,13 @@ from __future__ import annotations
 
 from typing import Any
 
-from artifex.generative_models.core.base import MLP
-from flax import nnx
 import jax
 import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from flax import nnx
 from jaxtyping import Array, Float
 
 from diffbio.losses.statistical_losses import zinb_negative_log_likelihood
-from diffbio.utils.nn_utils import ensure_rngs
 
 
 class CountVAEBackboneMixin:
@@ -32,10 +31,9 @@ class CountVAEBackboneMixin:
         latent_dim: int,
         hidden_dims: list[int],
         n_outputs: int,
-        rngs: nnx.Rngs | None,
+        rngs: nnx.Rngs,
     ) -> None:
         """Initialise the shared count-VAE encoder and decoder layers."""
-        safe_rngs = ensure_rngs(rngs)
 
         encoder_hidden_dims = list(hidden_dims)
         decoder_hidden_dims = list(reversed(hidden_dims))
@@ -47,7 +45,7 @@ class CountVAEBackboneMixin:
                 activation="relu",
                 output_activation="relu",
                 use_batch_norm=False,
-                rngs=safe_rngs,
+                rngs=rngs,
             )
             encoder_out_dim = encoder_hidden_dims[-1]
         else:
@@ -57,12 +55,12 @@ class CountVAEBackboneMixin:
         self.fc_mean = nnx.Linear(
             in_features=encoder_out_dim,
             out_features=latent_dim,
-            rngs=safe_rngs,
+            rngs=rngs,
         )
         self.fc_logvar = nnx.Linear(
             in_features=encoder_out_dim,
             out_features=latent_dim,
-            rngs=safe_rngs,
+            rngs=rngs,
         )
 
         if decoder_hidden_dims:
@@ -72,7 +70,7 @@ class CountVAEBackboneMixin:
                 activation="relu",
                 output_activation="relu",
                 use_batch_norm=False,
-                rngs=safe_rngs,
+                rngs=rngs,
             )
             decoder_out_dim = decoder_hidden_dims[-1]
         else:
@@ -82,17 +80,16 @@ class CountVAEBackboneMixin:
         self.fc_output = nnx.Linear(
             in_features=decoder_out_dim,
             out_features=n_outputs,
-            rngs=safe_rngs,
+            rngs=rngs,
         )
 
     def _init_count_vae_operator(
         self,
         *,
         config: Any,
-        rngs: nnx.Rngs | None,
+        rngs: nnx.Rngs,
     ) -> nnx.Rngs:
-        """Initialise shared count-VAE operator state and return safe RNGs."""
-        safe_rngs = ensure_rngs(rngs)
+        """Initialise shared count-VAE operator state and return the rngs it used."""
         self.n_genes = config.n_genes
         self.stream_name = nnx.static(config.stream_name)
         self._init_count_vae_backbone(
@@ -100,9 +97,9 @@ class CountVAEBackboneMixin:
             latent_dim=config.latent_dim,
             hidden_dims=config.hidden_dims,
             n_outputs=config.n_genes,
-            rngs=safe_rngs,
+            rngs=rngs,
         )
-        return safe_rngs
+        return rngs
 
     def encode(
         self,
@@ -145,7 +142,7 @@ class CountVAEBackbone(CountVAEBackboneMixin, nnx.Module):
         latent_dim: int,
         hidden_dims: list[int],
         n_outputs: int,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
     ) -> None:
         """Initialise a standalone shared count-VAE backbone."""
         super().__init__()

{diffbio-0.1.4 → diffbio-0.1.6}/src/diffbio/operators/_loss_balancing.py

@@ -10,21 +10,18 @@ from flax import nnx
 from jaxtyping import Array, Float
 from opifex.core.physics.gradnorm import GradNormBalancer
 
-from diffbio.utils.nn_utils import ensure_rngs
-
 
 def combine_scalar_losses(
     losses: Mapping[str, Float[Array, ""]],
     *,
-    use_gradnorm: bool,
-    rngs: nnx.Rngs | None = None,
+    balancer: GradNormBalancer | None,
 ) -> Float[Array, ""]:
-    """Combine scalar losses with optional GradNorm-based balancing.
+    """Combine scalar losses, weighted by ``balancer`` when one is given.
 
     Args:
         losses: Named scalar losses to combine.
-        use_gradnorm: Whether to balance losses with ``GradNormBalancer``.
-        rngs: Optional random generators used when constructing GradNorm.
+        balancer: The ``GradNormBalancer`` whose weights combine the losses, or ``None``
+            to sum them.
 
     Returns:
         Combined scalar loss.
@@ -35,15 +32,9 @@ def combine_scalar_losses(
     if not losses:
         msg = "losses must contain at least one scalar loss"
         raise ValueError(msg)
-
     loss_values = list(losses.values())
-    if use_gradnorm:
-        balancer = GradNormBalancer(
-            num_losses=len(loss_values),
-            rngs=ensure_rngs(rngs),
-        )
+    if balancer is not None:
         return balancer.compute_weighted_loss(jnp.stack(loss_values))
-
     total_loss = loss_values[0]
     for loss_value in loss_values[1:]:
         total_loss = total_loss + loss_value
@@ -51,16 +42,25 @@ def combine_scalar_losses(
 
 
 class LossBalancingMixin:
-    """Reusable operator mixin exposing ``compute_balanced_loss``."""
+    """Reusable operator mixin exposing ``compute_balanced_loss``.
+
+    The mixin is a stateless combiner: with ``config.use_gradnorm`` it builds a fresh
+    ``GradNormBalancer`` from the operator's ``rngs`` on every call and never updates its
+    weights, so it weights the losses equally; a training loop that wants adaptive GradNorm
+    composes the balancer itself across steps, as ``diffbio.pipelines.joint_training`` does.
+    """
 
     config: Any
+    rngs: nnx.Rngs
 
     def compute_balanced_loss(
         self,
         losses: Mapping[str, Float[Array, ""]],
     ) -> Float[Array, ""]:
         """Combine operator loss terms using the config's GradNorm flag."""
-        return combine_scalar_losses(
-            losses,
-            use_gradnorm=bool(getattr(self.config, "use_gradnorm", False)),
+        balancer = (
+            GradNormBalancer(num_losses=len(losses), rngs=self.rngs)
+            if getattr(self.config, "use_gradnorm", False)
+            else None
         )
+        return combine_scalar_losses(losses, balancer=balancer)