diffbio 0.1.5__tar.gz → 0.1.7__tar.gz

{diffbio-0.1.5 → diffbio-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.5
+Version: 0.1.7
 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.8
 Requires-Dist: chex>=0.1.7
-Requires-Dist: datarax>=0.1.10
+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'
@@ -352,10 +352,11 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 ### Training a Pipeline
 
 ```python
+from calibrax.metrics.functional import softmax_cross_entropy
+from substrax.optim import OptimizerConfig
 from diffbio.utils import (
     Trainer,
     TrainingConfig,
-    cross_entropy_loss,
     create_synthetic_training_data,
     data_iterator,
 )
@@ -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
@@ -382,10 +384,9 @@ trainer = Trainer(pipeline, config)
 
 # Define loss function
 def loss_fn(predictions, targets):
-    return cross_entropy_loss(
+    return softmax_cross_entropy(
         predictions["logits"],
         targets["labels"],
-        num_classes=3,
     )
 
 # Train

{diffbio-0.1.5 → diffbio-0.1.7}/README.md

@@ -138,10 +138,11 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 ### Training a Pipeline
 
 ```python
+from calibrax.metrics.functional import softmax_cross_entropy
+from substrax.optim import OptimizerConfig
 from diffbio.utils import (
     Trainer,
     TrainingConfig,
-    cross_entropy_loss,
     create_synthetic_training_data,
     data_iterator,
 )
@@ -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
@@ -168,10 +170,9 @@ trainer = Trainer(pipeline, config)
 
 # Define loss function
 def loss_fn(predictions, targets):
-    return cross_entropy_loss(
+    return softmax_cross_entropy(
         predictions["logits"],
         targets["labels"],
-        num_classes=3,
     )
 
 # Train

{diffbio-0.1.5 → diffbio-0.1.7}/pyproject.toml

@@ -37,15 +37,16 @@ classifiers = [
 ]
 dependencies = [
   # Datarax - operator, source, and pipeline contracts
-  "datarax>=0.1.10",
+  "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.8",
+  # 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.5"
+version = "0.1.7"
 
 [project.optional-dependencies]
 all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]

{diffbio-0.1.5 → diffbio-0.1.7}/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",
@@ -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,7 +274,7 @@ 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,
@@ -298,7 +300,9 @@ class EncoderDecoderOperator(OperatorModule):
             Sampled latent representation.
         """
         if key is None:
-            key = get_rng_key(self.rngs, "sample", fallback_seed=0)
+            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
@@ -473,7 +477,7 @@ class HMMOperator(OperatorModule):
         self,
         config: OperatorConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize HMMOperator.
@@ -490,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)
 

{diffbio-0.1.5 → diffbio-0.1.7}/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.5 → diffbio-0.1.7}/src/diffbio/losses/self_supervised_losses.py

@@ -11,10 +11,11 @@ term, which the ``core.losses`` ledger flags as the one genuinely new B6 piece.
 from __future__ import annotations
 
 import jax
+from calibrax.metrics.functional import mae, mse
 from flax import nnx
 
-_EPSILON = 1.0e-8
-_VALID_LOSS_TYPES = ("mse", "mae")
+_LOSSES = {"mse": mse, "mae": mae}
+_VALID_LOSS_TYPES = tuple(_LOSSES)
 
 
 def masked_value_loss(
@@ -23,28 +24,28 @@ def masked_value_loss(
     mask: jax.Array,
     *,
     loss_type: str = "mse",
-    epsilon: float = _EPSILON,
 ) -> jax.Array:
     """Reconstruction loss averaged over masked positions only.
 
+    The reduction is calibrax's masked mean: positions where ``mask`` is nonzero are
+    averaged, the rest contribute nothing, and an all-zero mask gives ``0`` rather than
+    ``0 / 0``.
+
     Args:
         predictions: Predicted values, any shape broadcastable with ``mask``.
         targets: Ground-truth values, same shape as ``predictions``.
         mask: ``1`` at positions to reconstruct (the masked-out values), ``0`` elsewhere.
         loss_type: ``"mse"`` (squared error) or ``"mae"`` (absolute error).
-        epsilon: Stabilizer so an all-zero mask returns ``0`` rather than ``0 / 0``.
 
     Returns:
-        The mask-weighted mean reconstruction error at the masked positions.
+        The mean reconstruction error at the masked positions.
 
     Raises:
         ValueError: If ``loss_type`` is not one of ``("mse", "mae")``.
     """
-    if loss_type not in _VALID_LOSS_TYPES:
+    if loss_type not in _LOSSES:
         raise ValueError(f"loss_type must be one of {_VALID_LOSS_TYPES}, got {loss_type!r}")
-    residual = predictions - targets
-    error = residual**2 if loss_type == "mse" else abs(residual)
-    return (mask * error).sum() / (mask.sum() + epsilon)
+    return _LOSSES[loss_type](predictions, targets, mask=mask)
 
 
 class MaskedValueLoss(nnx.Module):

{diffbio-0.1.5 → diffbio-0.1.7}/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.5 → diffbio-0.1.7}/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.5 → diffbio-0.1.7}/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)

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/alignment/profile_hmm.py

@@ -88,7 +88,7 @@ class ProfileHMMSearch(TemperatureOperator):
         self,
         config: ProfileHMMConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize the profile HMM operator.
@@ -100,9 +100,6 @@ class ProfileHMMSearch(TemperatureOperator):
         """
         super().__init__(config, rngs=rngs, name=name)
 
-        if rngs is None:
-            rngs = nnx.Rngs(0)
-
         self.profile_length = config.profile_length
         self.alphabet_size = config.alphabet_size
         # Temperature is now managed by TemperatureOperator via self._temperature

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/alignment/soft_msa.py

@@ -32,9 +32,9 @@ from diffbio.core.base_operators import TemperatureOperator
 from diffbio.utils.nn_utils import (
     ARTIFEX_GELU_MLP_KWARGS,
     ARTIFEX_GELU_NO_OUTPUT_MLP_KWARGS,
-    ensure_rngs,
 )
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -220,7 +220,7 @@ class SoftProgressiveMSA(TemperatureOperator):
         self,
         config: SoftProgressiveMSAConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize the soft progressive MSA operator.
@@ -232,8 +232,6 @@ class SoftProgressiveMSA(TemperatureOperator):
         """
         super().__init__(config, rngs=rngs, name=name)
 
-        rngs = ensure_rngs(rngs)
-
         self.hidden_dim = config.hidden_dim
         # Temperature is now managed by TemperatureOperator via self._temperature
         self.alphabet_size = config.alphabet_size

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/assembly/gnn_assembly.py

@@ -97,7 +97,7 @@ class GNNAssemblyNavigator(GraphOperator):
         self,
         config: GNNAssemblyNavigatorConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize the GNN assembly navigator.
@@ -109,9 +109,6 @@ class GNNAssemblyNavigator(GraphOperator):
         """
         super().__init__(config, rngs=rngs, name=name)
 
-        if rngs is None:
-            rngs = nnx.Rngs(0)
-
         self.hidden_dim = config.hidden_dim
 
         # Temperature management (similar to TemperatureOperator pattern)

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/drug_discovery/_graph_utils.py

@@ -7,7 +7,7 @@ import jax.numpy as jnp
 from flax import nnx
 
 from diffbio.operators.drug_discovery.message_passing import StackedMessagePassing
-from diffbio.utils.nn_utils import ensure_rngs
+
 
 logger = logging.getLogger(__name__)
 
@@ -44,7 +44,7 @@ def _require_config_attr(config: Any, attr: str) -> Any:
 def initialize_graph_encoder(
     module: Any,
     *,
-    rngs: nnx.Rngs | None,
+    rngs: nnx.Rngs,
     hidden_dim: int,
     num_layers: int,
     in_features: int,
@@ -52,7 +52,6 @@ def initialize_graph_encoder(
     attr: str = "encoder",
 ) -> nnx.Rngs:
     """Ensure RNGs and attach a message-passing encoder."""
-    resolved_rngs = ensure_rngs(rngs)
     setattr(
         module,
         attr,
@@ -60,18 +59,18 @@ def initialize_graph_encoder(
             hidden_dim=hidden_dim,
             num_layers=num_layers,
             in_features=in_features,
-            rngs=resolved_rngs,
+            rngs=rngs,
             num_edge_features=num_edge_features,
         ),
     )
-    return resolved_rngs
+    return rngs
 
 
 def initialize_graph_encoder_from_config(
     module: Any,
     config: Any,
     *,
-    rngs: nnx.Rngs | None,
+    rngs: nnx.Rngs,
     num_layers_attr: str = "num_message_passing_steps",
     hidden_dim_attr: str = "hidden_dim",
     in_features_attr: str = "in_features",

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/drug_discovery/admet_predictor.py

@@ -161,7 +161,7 @@ class ADMETPredictor(OperatorModule):
         self,
         config: ADMETConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
         name: str | None = None,
     ):
         """Initialize ADMET predictor.

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/drug_discovery/attentive_fp.py

@@ -221,7 +221,7 @@ class AttentiveFP(OperatorModule):
         self,
         config: AttentiveFPConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
     ):
         """Initialize AttentiveFP.
 
@@ -231,9 +231,6 @@ class AttentiveFP(OperatorModule):
         """
         super().__init__(config, rngs=rngs)
 
-        if rngs is None:
-            rngs = nnx.Rngs(0)
-
         # Initial linear projection
         self.input_proj = nnx.Linear(config.in_features, config.hidden_dim, rngs=rngs)
 

{diffbio-0.1.5 → diffbio-0.1.7}/src/diffbio/operators/drug_discovery/dti.py

@@ -105,12 +105,10 @@ class DifferentiableDTIPipeline(nnx.Module):
         self,
         config: DTIPipelineConfig,
         *,
-        rngs: nnx.Rngs | None = None,
+        rngs: nnx.Rngs,
     ) -> None:
         """Initialize protein encoder, drug encoder, and pair scorer."""
         super().__init__()
-        if rngs is None:
-            rngs = nnx.Rngs(0)
 
         self.config = nnx.static(config)
         self.protein_encoder = TransformerSequenceEncoder(