diffbio 0.1.4__tar.gz → 0.1.5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.4
+Version: 0.1.5
 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
@@ -50,7 +50,7 @@ Requires-Dist: beartype>=0.14.1
 Requires-Dist: biopython>=1.81
 Requires-Dist: calibrax>=0.1.5
 Requires-Dist: chex>=0.1.7
-Requires-Dist: datarax>=0.1.7
+Requires-Dist: datarax>=0.1.10
 Requires-Dist: flax>=0.12.9
 Requires-Dist: h5py>=3.7
 Requires-Dist: jax-md>=0.2.27

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

@@ -37,7 +37,7 @@ classifiers = [
 ]
 dependencies = [
   # Datarax - operator, source, and pipeline contracts
-  "datarax>=0.1.7",
+  "datarax>=0.1.10",
   # Artifex - modeling, modality, and transformer substrate
   "avitai-artifex>=0.1.5",
   # Opifex - scientific ML, operator learning, and advanced optimization
@@ -70,7 +70,7 @@ license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.4"
+version = "0.1.5"
 
 [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.5}/src/diffbio/core/base_operators.py

@@ -139,7 +139,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 +233,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."""
@@ -278,6 +278,7 @@ class EncoderDecoderOperator(OperatorModule):
         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 +289,16 @@ 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 = get_rng_key(self.rngs, "sample", fallback_seed=0)
         std = jnp.exp(0.5 * log_var)
         epsilon = jax.random.normal(key, mean.shape)
         return mean + std * epsilon
@@ -345,7 +351,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 +448,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."""
@@ -593,7 +599,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.5}/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.5}/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.5}/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.5}/src/diffbio/operators/_masked_gene_transformer.py

@@ -8,7 +8,6 @@ construction and one mask/input preparation flow.
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
 
 import jax
 import jax.numpy as jnp
@@ -63,28 +62,34 @@ def build_masked_gene_transformer_encoder(
 
 def build_masked_gene_mask(
     *,
-    random_params: Any,
+    key: jax.Array | None,
     mask_ratio: float,
     n_genes: int,
 ) -> Array:
-    """Build a per-gene binary mask for masked-gene transformer operators."""
-    if random_params is not None and mask_ratio > 0:
-        noise = jax.random.uniform(random_params, (n_genes,))
-        return (noise < mask_ratio).astype(jnp.float32)
-    return jnp.zeros(n_genes, dtype=jnp.float32)
+    """Build a per-gene binary mask for masked-gene transformer operators.
+
+    With ``mask_ratio == 0`` nothing is masked and no key is needed; otherwise the mask is
+    drawn from ``key``, which must be given.
+    """
+    if mask_ratio <= 0:
+        return jnp.zeros(n_genes, dtype=jnp.float32)
+    if key is None:
+        raise ValueError("masking genes needs the record's PRNG key")
+    noise = jax.random.uniform(key, (n_genes,))
+    return (noise < mask_ratio).astype(jnp.float32)
 
 
 def prepare_masked_gene_batch(
     data: PyTree,
     *,
-    random_params: Any,
+    key: jax.Array | None,
     mask_ratio: float,
 ) -> tuple[Array, Array, Array]:
     """Extract counts, int32 gene IDs, and the shared masking vector."""
     counts = data["counts"]
     gene_ids = jnp.asarray(data["gene_ids"], dtype=jnp.int32)
     mask = build_masked_gene_mask(
-        random_params=random_params,
+        key=key,
         mask_ratio=mask_ratio,
         n_genes=int(counts.shape[1]),
     )
@@ -96,23 +101,14 @@ class MaskedGeneTransformerOperatorMixin:
 
     config: MaskedGeneTransformerConfigBase
 
-    def generate_random_params(
-        self,
-        rng: jax.Array,
-        data_shapes: PyTree,
-    ) -> jax.Array:
-        """Return the RNG key used for reproducible masking inside apply."""
-        del data_shapes
-        return rng
-
     def prepare_masked_gene_batch(
         self,
         data: PyTree,
-        random_params: Any,
+        key: jax.Array | None,
     ) -> tuple[Array, Array, Array]:
         """Prepare shared masked-gene inputs for per-cell `vmap` execution."""
         return prepare_masked_gene_batch(
             data,
-            random_params=random_params,
+            key=key,
             mask_ratio=self.config.mask_ratio,
         )

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

@@ -41,7 +41,6 @@ class ProfileHMMConfig(OperatorConfig):
         learnable_profile: Whether profile parameters are learnable.
     """
 
-    cacheable: bool = True
     profile_length: int = 100
     alphabet_size: int = 20  # Amino acids by default
     temperature: float = 1.0
@@ -309,7 +308,7 @@ class ProfileHMMSearch(TemperatureOperator):
         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]:
         """Apply profile HMM search to sequence.
@@ -319,7 +318,7 @@ class ProfileHMMSearch(TemperatureOperator):
                 - "sequence": One-hot encoded sequence (seq_len, alphabet_size)
             state: Element state (passed through unchanged)
             metadata: Element metadata (passed through unchanged)
-            random_params: Not used (deterministic operator)
+            key: Unused.
             stats: Not used
 
         Returns:

{diffbio-0.1.4 → diffbio-0.1.5}/src/diffbio/operators/alignment/smith_waterman.py

@@ -39,7 +39,6 @@ class SmithWatermanConfig(TemperatureConfig):
         gap_extend: Penalty for extending a gap.
     """
 
-    cacheable: bool = True
     gap_open: float = DEFAULT_GAP_OPEN
     gap_extend: float = DEFAULT_GAP_EXTEND
 
@@ -216,7 +215,7 @@ class SmoothSmithWaterman(TemperatureOperator):
         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]:
         """Apply alignment to sequence pair data.
@@ -233,7 +232,7 @@ class SmoothSmithWaterman(TemperatureOperator):
                 - "seq2": Second sequence, one-hot encoded (len2, alphabet_size)
             state: Element state (passed through unchanged)
             metadata: Element metadata (passed through unchanged)
-            random_params: Not used (deterministic operator)
+            key: Unused.
             stats: Not used
 
         Returns:

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

@@ -339,7 +339,7 @@ class SoftProgressiveMSA(TemperatureOperator):
         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]:
         """Apply soft progressive MSA.
@@ -349,7 +349,7 @@ class SoftProgressiveMSA(TemperatureOperator):
                 - "sequences": Input sequences (n_seqs, seq_len, alphabet_size)
             state: Element state (passed through unchanged)
             metadata: Element metadata (passed through unchanged)
-            random_params: Not used
+            key: Unused.
             stats: Not used
 
         Returns:

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

@@ -20,6 +20,7 @@ import logging
 from dataclasses import dataclass
 from typing import Any
 
+import jax
 import jax.numpy as jnp
 from flax import nnx
 from jaxtyping import Array, Float, Int, PyTree
@@ -187,7 +188,7 @@ class GNNAssemblyNavigator(GraphOperator):
         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]:
         """Apply GNN assembly navigation.
@@ -199,7 +200,7 @@ class GNNAssemblyNavigator(GraphOperator):
                 - "edge_features": Edge features (n_edges, edge_features)
             state: Element state (passed through unchanged)
             metadata: Element metadata (passed through unchanged)
-            random_params: Not used
+            key: Unused.
             stats: Not used
 
         Returns:

{diffbio-0.1.4 → diffbio-0.1.5}/src/diffbio/operators/assembly/metagenomic_binning.py

@@ -15,6 +15,7 @@ import jax
 import jax.numpy as jnp
 from artifex.generative_models.core.base import MLP
 from flax import nnx
+from datarax.core.operator import require_key
 from jaxtyping import Array, Float
 
 from diffbio.configs import TemperatureConfig, apply_stochastic_sampling_defaults
@@ -212,7 +213,7 @@ class DifferentiableMetagenomicBinner(TemperatureOperator, EncoderDecoderOperato
         data: dict[str, Array],
         state: dict[str, Any],
         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[dict[str, Array], dict[str, Any], dict[str, Any] | None]:
         """Apply metagenomic binning.
@@ -223,7 +224,8 @@ class DifferentiableMetagenomicBinner(TemperatureOperator, EncoderDecoderOperato
                 - abundance: Float[Array, "n_contigs n_samples"]
             state: Element state (passed through).
             metadata: Element metadata (passed through).
-            random_params: Random parameters.
+            key: The record's PRNG key; in train mode the latent sample is drawn
+                from it, in eval mode the mean is used and no key is needed.
             stats: Optional statistics dict.
 
         Returns:
@@ -242,7 +244,7 @@ class DifferentiableMetagenomicBinner(TemperatureOperator, EncoderDecoderOperato
         if self.latent_sampling_mode.deterministic:
             z = mu  # Eval mode: deterministic
         else:
-            z = self.reparameterize(mu, logvar)  # Train mode: stochastic
+            z = self.reparameterize(mu, logvar, require_key(key, self))  # Train mode
 
         # Decode
         tnf_recon, abundance_recon = self.decode(z)

{diffbio-0.1.4 → diffbio-0.1.5}/src/diffbio/operators/crispr/guide_scoring.py

@@ -28,6 +28,7 @@ import logging
 from dataclasses import dataclass
 from typing import Any
 
+import jax
 import jax.numpy as jnp
 from artifex.generative_models.core.base import MLP
 from datarax.core.config import OperatorConfig
@@ -195,7 +196,7 @@ class DifferentiableCRISPRScorer(OperatorModule):
         data: dict[str, Any],
         state: dict[str, Any],
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Apply CRISPR scoring to guide sequences.
@@ -205,7 +206,7 @@ class DifferentiableCRISPRScorer(OperatorModule):
                 - "guides": One-hot encoded guides (n_guides, guide_length, 4).
             state: Per-element state (passed through).
             metadata: Optional metadata (passed through).
-            random_params: Random parameters for stochastic operations.
+            key: Unused.
             stats: Optional statistics dictionary.
 
         Returns:

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

@@ -12,11 +12,6 @@ from diffbio.utils.nn_utils import ensure_rngs
 logger = logging.getLogger(__name__)
 
 
-def stabilize_operator_id(module: Any) -> None:
-    """Mark operator unique ID as static for NNX/JAX transformations."""
-    module._unique_id = nnx.static(module._unique_id)
-
-
 def build_encoder(
     *,
     hidden_dim: int,
@@ -56,8 +51,7 @@ def initialize_graph_encoder(
     num_edge_features: int | None = None,
     attr: str = "encoder",
 ) -> nnx.Rngs:
-    """Stabilize ID, ensure RNGs, and attach a message-passing encoder."""
-    stabilize_operator_id(module)
+    """Ensure RNGs and attach a message-passing encoder."""
     resolved_rngs = ensure_rngs(rngs)
     setattr(
         module,

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

@@ -16,6 +16,7 @@ import logging
 from dataclasses import dataclass
 from typing import Any
 
+import jax
 import jax.numpy as jnp
 from artifex.generative_models.core.base import MLP
 from datarax.core.config import OperatorConfig
@@ -201,7 +202,7 @@ class ADMETPredictor(OperatorModule):
         data: dict[str, Any],
         state: dict[str, Any],
         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[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Predict ADMET properties from molecular graph.
@@ -214,7 +215,7 @@ class ADMETPredictor(OperatorModule):
                 - node_mask: (num_nodes,) mask for valid nodes
             state: Per-element state (passed through).
             metadata: Optional metadata.
-            random_params: Unused random parameters.
+            key: Unused.
             stats: Optional statistics dictionary.
 
         Returns:
@@ -223,7 +224,7 @@ class ADMETPredictor(OperatorModule):
                 - unchanged state
                 - unchanged metadata
         """
-        del random_params, stats  # Unused
+        del key, stats  # Unused
 
         graph_repr = graph_sum_readout(data, self.encoder, dropout=self.dropout)
 

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

@@ -28,7 +28,6 @@ from flax import nnx
 
 from diffbio.operators.drug_discovery._graph_utils import (
     build_optional_dropout,
-    stabilize_operator_id,
 )
 
 logger = logging.getLogger(__name__)
@@ -232,8 +231,6 @@ class AttentiveFP(OperatorModule):
         """
         super().__init__(config, rngs=rngs)
 
-        stabilize_operator_id(self)
-
         if rngs is None:
             rngs = nnx.Rngs(0)
 
@@ -292,7 +289,7 @@ class AttentiveFP(OperatorModule):
         data: dict[str, Any],
         state: dict[str, Any],
         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[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Compute AttentiveFP molecular fingerprint.
@@ -305,7 +302,7 @@ class AttentiveFP(OperatorModule):
                 - node_mask: (num_nodes,) optional mask for valid nodes
             state: Per-element state (passed through).
             metadata: Optional metadata.
-            random_params: Unused random parameters.
+            key: Unused.
             stats: Optional statistics dictionary.
 
         Returns:
@@ -314,7 +311,7 @@ class AttentiveFP(OperatorModule):
                 - unchanged state
                 - unchanged metadata
         """
-        del random_params, stats  # Unused
+        del key, stats  # Unused
 
         node_features = data["node_features"]
         adjacency = data["adjacency"]

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

@@ -14,6 +14,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 from datarax.core.config import OperatorConfig
@@ -150,11 +151,11 @@ class DifferentiableDTIPipeline(nnx.Module):
         data: dict[str, Any],
         state: dict[str, Any],
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Score a prepared paired protein/drug DTI batch."""
-        del random_params, stats
+        del key, stats
 
         protein_result, _, _ = self.protein_encoder.apply(
             {

{diffbio-0.1.4 → diffbio-0.1.5}/src/diffbio/operators/drug_discovery/fingerprint.py

@@ -12,6 +12,7 @@ import logging
 from dataclasses import dataclass
 from typing import Any
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 from datarax.core.config import OperatorConfig
@@ -25,7 +26,6 @@ from diffbio.operators.drug_discovery._graph_utils import (
     ensure_rngs,
     graph_sum_readout,
     initialize_graph_encoder,
-    stabilize_operator_id,
     unpack_graph_inputs,
 )
 
@@ -110,7 +110,7 @@ class DifferentiableMolecularFingerprint(OperatorModule):
         data: dict[str, Any],
         state: dict[str, Any],
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Compute molecular fingerprint.
@@ -122,7 +122,7 @@ class DifferentiableMolecularFingerprint(OperatorModule):
                 - node_mask: (num_nodes,) mask for valid nodes
             state: Per-element state (passed through).
             metadata: Optional metadata.
-            random_params: Unused random parameters.
+            key: Unused.
             stats: Optional statistics dictionary.
 
         Returns:
@@ -244,7 +244,6 @@ class CircularFingerprintOperator(OperatorModule):
         """
         super().__init__(config, rngs=rngs)
 
-        stabilize_operator_id(self)
         rngs = ensure_rngs(rngs)
 
         if config.differentiable:
@@ -368,7 +367,7 @@ class CircularFingerprintOperator(OperatorModule):
         data: dict[str, Any],
         state: dict[str, Any],
         metadata: dict[str, Any] | None,
-        random_params: Any = None,
+        key: jax.Array | None = None,
         stats: dict[str, Any] | None = None,
     ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
         """Compute circular fingerprint.
@@ -383,7 +382,7 @@ class CircularFingerprintOperator(OperatorModule):
                     - smiles: SMILES string
             state: Per-element state (passed through).
             metadata: Optional metadata.
-            random_params: Unused random parameters.
+            key: Unused.
             stats: Optional statistics dictionary.
 
         Returns: