diffbio 0.1.3__tar.gz → 0.1.5__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.3
-Summary: DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Artifex, Opifex, and Calibrax
+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
 Project-URL: Source, https://github.com/avitai/DiffBio
@@ -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
@@ -74,6 +74,7 @@ Requires-Dist: griffe>=1.7.3; extra == 'all'
 Requires-Dist: import-linter>=2.5; extra == 'all'
 Requires-Dist: interrogate>=1.7.0; extra == 'all'
 Requires-Dist: ipykernel>=6.29.5; extra == 'all'
+Requires-Dist: jax-metal>=0.1.0; (sys_platform == 'darwin' and platform_machine == 'arm64') and extra == 'all'
 Requires-Dist: jax[cuda12]>=0.11.1; extra == 'all'
 Requires-Dist: lineax>=0.0.8; extra == 'all'
 Requires-Dist: matplotlib>=3.7; extra == 'all'
@@ -101,6 +102,7 @@ Requires-Dist: pytest-timeout>=2.1; extra == 'all'
 Requires-Dist: pytest-xdist>=3.6; extra == 'all'
 Requires-Dist: pytest>=8.3.5; extra == 'all'
 Requires-Dist: python-dotenv>=1; extra == 'all'
+Requires-Dist: pyyaml>=6; extra == 'all'
 Requires-Dist: radon>=6.0.1; extra == 'all'
 Requires-Dist: rdkit>=2025.9.3; extra == 'all'
 Requires-Dist: ruff>=0.1.5; extra == 'all'
@@ -185,6 +187,8 @@ Requires-Dist: pymdown-extensions>=10.14.3; extra == 'docs'
 Provides-Extra: genomics
 Requires-Dist: pyfaidx>=0.8.0; extra == 'genomics'
 Requires-Dist: pysam>=0.22.0; extra == 'genomics'
+Provides-Extra: metal
+Requires-Dist: jax-metal>=0.1.0; (sys_platform == 'darwin' and platform_machine == 'arm64') and extra == 'metal'
 Provides-Extra: soft-ops-advanced
 Requires-Dist: lineax>=0.0.8; extra == 'soft-ops-advanced'
 Requires-Dist: optimistix>=0.0.9; extra == 'soft-ops-advanced'
@@ -203,6 +207,7 @@ Requires-Dist: pytest-randomly>=3.16.0; extra == 'test'
 Requires-Dist: pytest-timeout>=2.1; extra == 'test'
 Requires-Dist: pytest-xdist>=3.6; extra == 'test'
 Requires-Dist: pytest>=8.3.5; extra == 'test'
+Requires-Dist: pyyaml>=6; extra == 'test'
 Provides-Extra: torch-io
 Requires-Dist: torch>=1.13.0; extra == 'torch-io'
 Description-Content-Type: text/markdown

{diffbio-0.1.3 → diffbio-0.1.5}/benchmarks/README.md

@@ -90,13 +90,13 @@ Download datasets before running:
 
 ```bash
 # Single-cell (immune_human, 2GB)
-mkdir -p /media/mahdi/ssd23/Data/scib
-wget -O /media/mahdi/ssd23/Data/scib/Immune_ALL_human.h5ad \
+mkdir -p /mnt/ssd2/Data/scib
+wget -O /mnt/ssd2/Data/scib/Immune_ALL_human.h5ad \
     "https://ndownloader.figshare.com/files/25717328"
 
 # Trajectory (pancreas, 51MB)
-mkdir -p /media/mahdi/ssd23/Data/scvelo
-wget -O /media/mahdi/ssd23/Data/scvelo/endocrinogenesis_day15.h5ad \
+mkdir -p /mnt/ssd2/Data/scvelo
+wget -O /mnt/ssd2/Data/scvelo/endocrinogenesis_day15.h5ad \
     "https://github.com/theislab/scvelo_notebooks/raw/master/data/Pancreas/endocrinogenesis_day15.h5ad"
 ```
 

{diffbio-0.1.3 → 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
@@ -64,16 +64,16 @@ dependencies = [
   "h5py>=3.7",
   "jax-md>=0.2.27"
 ]
-description = "DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Artifex, Opifex, and Calibrax"
+description = "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"
 keywords = ["jax", "flax", "bioinformatics", "differentiable", "variant-calling", "alignment", "machine-learning"]
 license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.3"
+version = "0.1.5"
 
 [project.optional-dependencies]
-all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
+all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
 benchmark = [
   # benchmarks/crossmodality/*.py call datasets.load_dataset. It used to arrive
   # transitively through the sibling packages; they now declare their heavy
@@ -131,6 +131,8 @@ docs = [
   "pymdown-extensions>=10.14.3"
 ]
 genomics = ["pysam>=0.22.0", "pyfaidx>=0.8.0"]
+# Apple Silicon GPU backend; setup.sh selects it on arm64 macOS.
+metal = ["jax-metal>=0.1.0; sys_platform == 'darwin' and platform_machine == 'arm64'"]
 soft-ops-advanced = ["optimistix>=0.0.9", "lineax>=0.0.8"]
 soft-ops-ot = ["diffbio[soft-ops-advanced]", "ott-jax>=0.5.0"]
 test = [
@@ -143,7 +145,9 @@ test = [
   "pytest-env>=1.0.1",
   "pytest-randomly>=3.16.0",
   "pytest-timeout>=2.1",
-  "pytest-xdist>=3.6"
+  "pytest-xdist>=3.6",
+  # tests/test_ci_shards.py reads the CI workflow.
+  "pyyaml>=6"
 ]
 torch-io = ["torch>=1.13.0"]
 
@@ -341,7 +345,7 @@ save_data = true
 warmup = true
 warmup_iterations = 3
 
-[tool.pytest-env]
+[tool.pytest_env]
 JAX_ENABLE_X64 = "0"
 JAX_SKIP_CUDA_CONSTRAINTS_CHECK = "1"
 LD_LIBRARY_PATH = "/usr/local/cuda/lib64"

{diffbio-0.1.3 → 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.3 → 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.3 → diffbio-0.1.5}/src/diffbio/core/soft_ops/__init__.py

@@ -81,6 +81,9 @@ from diffbio.core.soft_ops.sorting import (
     top_k_mask,
 )
 
+# --- Range-aware normalization ---
+from diffbio.core.soft_ops.normalization import temperature_softmax
+
 # --- Quantile ---
 from diffbio.core.soft_ops.quantile import (
     argmedian,
@@ -126,6 +129,7 @@ from diffbio.core.soft_ops.straight_through import (
 )
 
 __all__ = [
+    "temperature_softmax",
     # Types
     "SoftBool",
     "SoftIndex",

diffbio-0.1.5/src/diffbio/core/soft_ops/normalization.py

@@ -0,0 +1,174 @@
+"""Temperature softmax with jointly evaluated derivative coefficients.
+
+Ordinary division AD forms ``T**-2`` before multiplying the softmax tail.
+The product can be finite even when that reciprocal overflows or the tail
+underflows. Evaluate probability products, score gaps and reciprocal powers
+together in log space. A recursive coefficient JVP preserves mixed derivatives
+at zero gaps; differentiating an outer zero selection would lose them.
+
+This requires representable coefficient sums, not merely a representable final
+contraction. Work grows with donor count and derivative order: the first JVP
+has quadratic donor work. It is intended for small donor panels.
+"""
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from diffbio.core.soft_ops._utils import canonicalize_axis
+
+
+type _Specification = tuple[tuple[int, ...], tuple[tuple[int, int], ...], int]
+type _Inputs = tuple[Array, Array, Array]
+
+
+def _logits(scores: Array, temperature: Array, covered: Array) -> Array:
+    """Shift over covered donors before division, guarding excluded operands."""
+    maximum = jnp.max(jnp.where(covered, scores, -jnp.inf), axis=0, keepdims=True)
+    shifted = jnp.where(covered, scores - maximum, 0.0)
+    return jnp.where(covered, shifted / temperature, -jnp.inf)
+
+
+def _coefficient_value(
+    scores: Array, temperature: Array, covered: Array, specification: _Specification
+) -> Array:
+    """Evaluate a probability/gap monomial divided by a temperature power."""
+    indices, gaps, power = specification
+    log_weights = jax.nn.log_softmax(_logits(scores, temperature, covered), axis=0)
+    magnitude = sum(log_weights[index] for index in indices) - power * jnp.log(temperature)
+    sign = jnp.ones_like(magnitude)
+    nonzero = jnp.ones_like(magnitude, dtype=jnp.bool_)
+    for left, right in gaps:
+        gap = scores[left] - scores[right]
+        nonzero = nonzero & (gap != 0)
+        magnitude = magnitude + jnp.log(jnp.where(gap != 0, jnp.abs(gap), 1.0))
+        sign = sign * jnp.sign(gap)
+    # Guard the exponential too: a zero gap must not create 0 * inf.
+    value = sign * jnp.exp(jnp.where(nonzero, magnitude, 0.0))
+    return jnp.where(nonzero, value, 0.0)
+
+
+_coefficient = jax.custom_jvp(_coefficient_value, nondiff_argnums=(3,))
+
+
+def _coefficient_jvp(
+    specification: _Specification, primals: _Inputs, tangents: _Inputs
+) -> tuple[Array, Array]:
+    """Differentiate complete coefficients, including zero-gap extensions."""
+    scores, _, _ = primals
+    score_dot, temperature_dot, _ = tangents
+    indices, gaps, power = specification
+    value = _coefficient(*primals, specification)
+    derivative = jnp.zeros_like(value)
+    thermal = jnp.zeros_like(value)
+    for index in indices:
+        for other in range(scores.shape[0]):
+            if other != index:
+                extended = (*indices, other)
+                coefficient = _coefficient(*primals, (extended, gaps, power + 1))
+                derivative = derivative + coefficient * (score_dot[index] - score_dot[other])
+                thermal = thermal + _coefficient(
+                    *primals, (extended, (*gaps, (other, index)), power + 2)
+                )
+    for position, (left, right) in enumerate(gaps):
+        remaining = gaps[:position] + gaps[position + 1 :]
+        derivative = derivative + _coefficient(*primals, (indices, remaining, power)) * (
+            score_dot[left] - score_dot[right]
+        )
+    if power:
+        thermal = thermal - power * _coefficient(*primals, (indices, gaps, power + 1))
+    return value, derivative + thermal * temperature_dot
+
+
+_coefficient.defjvp(_coefficient_jvp)
+
+
+@jax.custom_jvp
+def _weights(scores: Array, temperature: Array, covered: Array) -> Array:
+    """Keep the native shifted forward softmax arithmetic."""
+    return jax.nn.softmax(_logits(scores, temperature, covered), axis=0)
+
+
+def _weights_jvp(primals: _Inputs, tangents: _Inputs) -> tuple[Array, Array]:
+    """Contract pairwise sensitivities without rounding either probability first."""
+    scores, _, _ = primals
+    score_dot, temperature_dot, _ = tangents
+    value = _weights(*primals)
+    rows = []
+    for index in range(scores.shape[0]):
+        derivative = jnp.zeros_like(value[index])
+        for other in range(scores.shape[0]):
+            if other != index:
+                probability = (index, other)
+                spatial = _coefficient(*primals, (probability, (), 1))
+                thermal = _coefficient(*primals, (probability, ((other, index),), 2))
+                derivative = derivative + spatial * (score_dot[index] - score_dot[other])
+                derivative = derivative + thermal * temperature_dot
+        rows.append(derivative)
+    return value, jnp.stack(rows)
+
+
+_weights.defjvp(_weights_jvp)
+
+
+@partial(jax.jit, static_argnames=("axis",))
+def temperature_softmax(
+    scores: Array,
+    temperature: float | Array = 1.0,
+    *,
+    axis: int = -1,
+    where: Array | None = None,
+) -> Array:
+    """Normalize small score axes with range-aware temperature derivatives.
+
+    Equivalent to masked ``jax.nn.softmax(scores / temperature, axis=axis)``
+    for valid inputs. Shift before division and evaluate complete derivative
+    coefficients in signed log space, preserving sensitivities even when a
+    probability rounds to zero or one. Supports JVP, VJP and higher derivatives.
+
+    This is an opt-in operation for small axes, not the default sorting kernel:
+    first derivatives require quadratic axis work and higher orders cost more.
+    Coefficients, score differences and their necessary sums must be representable;
+    a representable final contraction alone does not guarantee a finite derivative.
+
+    Args:
+        scores: Real floating array. Included scores and their pairwise differences
+            must be finite. Excluded scores are ignored, including NaN/infinity.
+        temperature: Positive finite real scalar. Callers own runtime validation;
+            nonpositive/nonfinite values are outside this numerical contract.
+        axis: Static reduction axis. Must be nonempty.
+        where: Boolean mask broadcastable to the score shape. Empty slices return
+            zero weights and zero derivatives, matching native masked softmax.
+
+    Returns:
+        Weights with the score shape and standard floating dtype promotion.
+
+    Raises:
+        TypeError: Scores are not floating, temperature is complex, or mask is not boolean.
+        ValueError: Temperature is not scalar, the axis is invalid/empty, or mask cannot broadcast.
+    """
+    scores = jnp.asarray(scores)
+    temperature = jnp.asarray(temperature)
+    if not jnp.issubdtype(scores.dtype, jnp.floating):
+        raise TypeError("scores must have a real floating dtype")
+    if jnp.issubdtype(temperature.dtype, jnp.complexfloating):
+        raise TypeError("temperature must be real")
+    if temperature.ndim != 0:
+        raise ValueError("temperature must be scalar")
+    temperature = temperature.astype(jnp.result_type(scores, temperature))
+    axis = canonicalize_axis(axis, scores.ndim)
+    if scores.shape[axis] == 0:
+        raise ValueError("the normalization axis must be nonempty")
+    covered = jnp.ones_like(scores, dtype=jnp.bool_) if where is None else jnp.asarray(where)
+    if covered.dtype != jnp.bool_:
+        raise TypeError("where must have boolean dtype")
+    covered = jnp.broadcast_to(covered, scores.shape)
+    scores, covered = jnp.moveaxis(scores, axis, 0), jnp.moveaxis(covered, axis, 0)
+    scores = jnp.where(covered, scores, 0.0)
+    populated = jnp.any(covered, axis=0, keepdims=True)
+    first = (jnp.arange(scores.shape[0]) == 0).reshape((-1,) + (1,) * (scores.ndim - 1))
+    safe_coverage = covered | (first & ~populated)
+    value = _weights(scores, temperature, safe_coverage)
+    return jnp.moveaxis(jnp.where(populated, value, 0.0), 0, axis)

{diffbio-0.1.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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.3 → 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: