diffbio 0.1.0__py3-none-any.whl

diffbio/operators/drug_discovery/admet_predictor.py

@@ -0,0 +1,285 @@
+"""Multi-task ADMET property prediction operator.
+
+This module implements a ChemProp-style multi-task ADMET predictor
+for predicting Absorption, Distribution, Metabolism, Excretion, and
+Toxicity properties of drug candidates.
+
+The implementation follows the TDC ADMET Benchmark with 22 standard endpoints.
+
+References:
+    - https://tdcommons.ai/benchmark/admet_group/overview/
+    - https://github.com/chemprop/chemprop
+    - Swanson et al. "ADMET-AI" Bioinformatics 2024
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+from diffbio.operators.drug_discovery._graph_utils import (
+    build_optional_dropout,
+    graph_sum_readout,
+    initialize_graph_encoder_from_config,
+)
+from diffbio.utils.nn_utils import ARTIFEX_RELU_MLP_KWARGS
+
+logger = logging.getLogger(__name__)
+
+# Standard TDC ADMET benchmark task names (22 tasks)
+ADMET_TASK_NAMES: list[str] = [
+    # Absorption (6)
+    "Caco2_Wang",
+    "HIA_Hou",
+    "Pgp_Broccatelli",
+    "Bioavailability_Ma",
+    "Lipophilicity_AstraZeneca",
+    "Solubility_AqSolDB",
+    # Distribution (3)
+    "BBB_Martins",
+    "PPBR_AZ",
+    "VDss_Lombardo",
+    # Metabolism (6)
+    "CYP2C9_Veith",
+    "CYP2D6_Veith",
+    "CYP3A4_Veith",
+    "CYP2C9_Substrate_CarbonMangels",
+    "CYP2D6_Substrate_CarbonMangels",
+    "CYP3A4_Substrate_CarbonMangels",
+    # Excretion (3)
+    "Half_Life_Obach",
+    "Clearance_Hepatocyte_AZ",
+    "Clearance_Microsome_AZ",
+    # Toxicity (4)
+    "LD50_Zhu",
+    "hERG",
+    "AMES",
+    "DILI",
+]
+
+# Task types: classification or regression
+ADMET_TASK_TYPES: dict[str, str] = {
+    # Absorption
+    "Caco2_Wang": "regression",
+    "HIA_Hou": "classification",
+    "Pgp_Broccatelli": "classification",
+    "Bioavailability_Ma": "classification",
+    "Lipophilicity_AstraZeneca": "regression",
+    "Solubility_AqSolDB": "regression",
+    # Distribution
+    "BBB_Martins": "classification",
+    "PPBR_AZ": "regression",
+    "VDss_Lombardo": "regression",
+    # Metabolism
+    "CYP2C9_Veith": "classification",
+    "CYP2D6_Veith": "classification",
+    "CYP3A4_Veith": "classification",
+    "CYP2C9_Substrate_CarbonMangels": "classification",
+    "CYP2D6_Substrate_CarbonMangels": "classification",
+    "CYP3A4_Substrate_CarbonMangels": "classification",
+    # Excretion
+    "Half_Life_Obach": "regression",
+    "Clearance_Hepatocyte_AZ": "regression",
+    "Clearance_Microsome_AZ": "regression",
+    # Toxicity
+    "LD50_Zhu": "regression",
+    "hERG": "classification",
+    "AMES": "classification",
+    "DILI": "classification",
+}
+
+
+@dataclass(frozen=True)
+class ADMETConfig(OperatorConfig):
+    # pylint: disable=too-many-instance-attributes
+    """Configuration for ADMET property predictor.
+
+    Attributes:
+        hidden_dim: Hidden dimension for message passing (default: 300).
+        num_message_passing_steps: Number of D-MPNN iterations (default: 3).
+        num_tasks: Number of ADMET prediction tasks (default: 22).
+        dropout_rate: Dropout rate for regularization (default: 0.0).
+        in_features: Number of input node features (default: 4).
+        num_edge_features: Number of edge features (default: 4).
+        ffn_hidden_dim: FFN hidden dimension (default: same as hidden_dim).
+        ffn_num_layers: Number of FFN layers (default: 2).
+        apply_task_activations: Apply sigmoid for classification tasks (default: False).
+    """
+
+    hidden_dim: int = 300
+    num_message_passing_steps: int = 3
+    num_tasks: int = 22
+    dropout_rate: float = 0.0
+    in_features: int = 4
+    num_edge_features: int = 4
+    ffn_hidden_dim: int | None = None
+    ffn_num_layers: int = 2
+    apply_task_activations: bool = False
+
+
+class ADMETPredictor(OperatorModule):
+    """Multi-task ADMET property predictor.
+
+    Implements a ChemProp-style directed message passing neural network
+    for predicting multiple ADMET properties simultaneously. The architecture
+    uses a shared molecular encoder with task-specific prediction heads.
+
+    Architecture:
+        1. Message passing encoder (D-MPNN style)
+        2. Graph-level readout via sum pooling
+        3. Shared feed-forward layers
+        4. Task-specific output heads
+
+    The 22 standard TDC ADMET endpoints cover:
+        - Absorption: Caco2, HIA, Pgp, Bioavailability, Lipophilicity, Solubility
+        - Distribution: BBB, PPBR, VDss
+        - Metabolism: CYP enzymes (2C9, 2D6, 3A4) inhibition and substrate
+        - Excretion: Half-life, Hepatocyte clearance, Microsome clearance
+        - Toxicity: LD50, hERG, AMES, DILI
+
+    Example:
+        ```python
+        config = ADMETConfig(hidden_dim=256, num_tasks=22)
+        predictor = ADMETPredictor(config, rngs=nnx.Rngs(42))
+        data = {"node_features": nodes, "adjacency": adj, "node_mask": mask}
+        result, _, _ = predictor.apply(data, {}, None)
+        predictions = result["predictions"]  # shape: (22,)
+        ```
+
+    References:
+        - https://tdcommons.ai/benchmark/admet_group/overview/
+        - Yang et al. "Analyzing Learned Molecular Representations" JCIM 2019
+    """
+
+    def __init__(
+        self,
+        config: ADMETConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ):
+        """Initialize ADMET predictor.
+
+        Args:
+            config: ADMET configuration.
+            rngs: Flax NNX random number generators.
+            name: Optional name for the operator.
+        """
+        super().__init__(config, rngs=rngs, name=name)
+
+        rngs = initialize_graph_encoder_from_config(self, config, rngs=rngs)
+
+        # FFN hidden dim defaults to hidden_dim
+        ffn_hidden = config.ffn_hidden_dim or config.hidden_dim
+
+        if config.ffn_num_layers > 1:
+            self.ffn_backbone = MLP(
+                hidden_dims=[ffn_hidden] * (config.ffn_num_layers - 1),
+                in_features=config.hidden_dim,
+                dropout_rate=config.dropout_rate,
+                rngs=rngs,
+                **ARTIFEX_RELU_MLP_KWARGS,
+            )
+        else:
+            self.ffn_backbone = None
+
+        # Task-specific output heads (one per ADMET task)
+        last_hidden = ffn_hidden if config.ffn_num_layers > 1 else config.hidden_dim
+        task_heads = [nnx.Linear(last_hidden, 1, rngs=rngs) for _ in range(config.num_tasks)]
+        self.task_heads = nnx.List(task_heads)
+
+        # Dropout
+        self.dropout = build_optional_dropout(config.dropout_rate, rngs=rngs)
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = 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.
+
+        Args:
+            data: Input data containing:
+                - node_features: (num_nodes, num_features) atom features
+                - adjacency: (num_nodes, num_nodes) adjacency matrix
+                - edge_features: Optional (num_nodes, num_nodes, num_edge_features)
+                - node_mask: (num_nodes,) mask for valid nodes
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with added "predictions" and "task_predictions" keys
+                - unchanged state
+                - unchanged metadata
+        """
+        del random_params, stats  # Unused
+
+        graph_repr = graph_sum_readout(data, self.encoder, dropout=self.dropout)
+
+        h = graph_repr
+        if self.ffn_backbone is not None:
+            ffn_output = self.ffn_backbone(h)
+            if isinstance(ffn_output, tuple):
+                raise TypeError("ADMETPredictor shared FFN must return a single tensor output.")
+            h = ffn_output
+
+        # Task-specific predictions
+        task_predictions = []
+        for i, head in enumerate(self.task_heads):
+            pred = head(h).squeeze(-1)
+
+            # Apply activation for classification tasks if configured
+            if self.config.apply_task_activations:
+                task_name = ADMET_TASK_NAMES[i] if i < len(ADMET_TASK_NAMES) else f"task_{i}"
+                if ADMET_TASK_TYPES.get(task_name) == "classification":
+                    pred = nnx.sigmoid(pred)
+
+            task_predictions.append(pred)
+
+        # Stack all predictions
+        predictions = jnp.stack(task_predictions)
+
+        result = {
+            **data,
+            "predictions": predictions,
+            "task_predictions": task_predictions,
+            "graph_representation": graph_repr,
+        }
+
+        return result, state, metadata
+
+
+def create_admet_predictor(
+    hidden_dim: int = 300,
+    num_layers: int = 3,
+    dropout_rate: float = 0.0,
+    seed: int = 42,
+) -> ADMETPredictor:
+    """Create an ADMET predictor with standard configuration.
+
+    Args:
+        hidden_dim: Hidden dimension for message passing.
+        num_layers: Number of message passing steps.
+        dropout_rate: Dropout rate.
+        seed: Random seed.
+
+    Returns:
+        Configured ADMETPredictor.
+    """
+    config = ADMETConfig(
+        hidden_dim=hidden_dim,
+        num_message_passing_steps=num_layers,
+        dropout_rate=dropout_rate,
+    )
+    return ADMETPredictor(config, rngs=nnx.Rngs(seed))

diffbio/operators/drug_discovery/attentive_fp.py

@@ -0,0 +1,411 @@
+"""AttentiveFP: Attention-based graph fingerprint for molecular property prediction.
+
+This module implements the AttentiveFP architecture from Xiong et al. 2019,
+which combines graph attention mechanisms with GRU cells for molecular
+representation learning.
+
+The architecture provides:
+
+    - Interpretable attention weights showing atom importance
+    - Two-level aggregation (atom-level and molecule-level)
+    - GRU-based state updates for iterative refinement
+
+References:
+    - Xiong et al. "Pushing the Boundaries of Molecular Representation for
+      Drug Discovery with the Graph Attention Mechanism" JCIM 2019
+    - https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.AttentiveFP.html
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+from diffbio.operators.drug_discovery._graph_utils import (
+    build_optional_dropout,
+    stabilize_operator_id,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _AttentiveFPArchitectureConfig:
+    """Architecture configuration for AttentiveFP."""
+
+    hidden_dim: int = 200
+    out_dim: int = 200
+    num_layers: int = 2
+    num_timesteps: int = 2
+    dropout_rate: float = 0.0
+
+
+@dataclass(frozen=True)
+class _AttentiveFPInputConfig:
+    """Graph input configuration for AttentiveFP."""
+
+    in_features: int = 39
+    edge_dim: int = 10
+    negative_slope: float = 0.2
+
+
+@dataclass(frozen=True)
+class AttentiveFPConfig(
+    _AttentiveFPArchitectureConfig,
+    _AttentiveFPInputConfig,
+    OperatorConfig,
+):
+    """Configuration for AttentiveFP operator."""
+
+    def __post_init__(self) -> None:
+        """Validate the AttentiveFP configuration."""
+        super().__post_init__()
+
+        if self.hidden_dim <= 0:
+            raise ValueError("hidden_dim must be positive.")
+        if self.out_dim <= 0:
+            raise ValueError("out_dim must be positive.")
+        if self.num_layers <= 0:
+            raise ValueError("num_layers must be positive.")
+        if self.num_timesteps <= 0:
+            raise ValueError("num_timesteps must be positive.")
+        if not 0.0 <= self.dropout_rate < 1.0:
+            raise ValueError("dropout_rate must be in [0.0, 1.0).")
+        if self.in_features <= 0:
+            raise ValueError("in_features must be positive.")
+        if self.edge_dim < 0:
+            raise ValueError("edge_dim must be non-negative.")
+        if self.negative_slope < 0.0:
+            raise ValueError("negative_slope must be non-negative.")
+
+
+class GATEConv(nnx.Module):
+    """Graph Attention with Edge features (GATE) convolution layer.
+
+    Combines node features with edge features using attention mechanism.
+    This is the core building block for AttentiveFP's atom-level processing.
+
+    The attention mechanism computes:
+        alpha_ij = softmax_j(LeakyReLU(a^T [Wh_i || Wh_j || We_ij]))
+
+    where || denotes concatenation and e_ij are edge features.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        edge_dim: int,
+        *,
+        negative_slope: float = 0.2,
+        rngs: nnx.Rngs,
+    ):
+        """Initialize GATE convolution.
+
+        Args:
+            in_dim: Input feature dimension.
+            out_dim: Output feature dimension.
+            edge_dim: Edge feature dimension.
+            negative_slope: LeakyReLU negative slope.
+            rngs: Random number generators.
+        """
+        super().__init__()
+        self.negative_slope = negative_slope
+
+        # Linear transformations
+        self.linear_src = nnx.Linear(in_dim, out_dim, rngs=rngs)
+        self.linear_dst = nnx.Linear(in_dim, out_dim, rngs=rngs)
+
+        # Edge feature projection (if edge_dim > 0)
+        self.use_edge_features = edge_dim > 0
+        if self.use_edge_features:
+            self.linear_edge = nnx.Linear(edge_dim, out_dim, rngs=rngs)
+
+        # Attention coefficients
+        # Attention is computed as: a^T [src || dst || edge]
+        attn_dim = out_dim * 3 if self.use_edge_features else out_dim * 2
+        self.attn = nnx.Linear(attn_dim, 1, use_bias=False, rngs=rngs)
+
+    def __call__(
+        self,
+        node_features: jnp.ndarray,
+        adjacency: jnp.ndarray,
+        edge_features: jnp.ndarray | None = None,
+    ) -> tuple[jnp.ndarray, jnp.ndarray]:
+        """Apply GATE convolution.
+
+        Args:
+            node_features: (num_nodes, in_dim) node features
+            adjacency: (num_nodes, num_nodes) adjacency matrix
+            edge_features: Optional (num_nodes, num_nodes, edge_dim)
+
+        Returns:
+            Tuple of (updated_features, attention_weights)
+        """
+        num_nodes = node_features.shape[0]
+
+        # Transform source and destination features
+        h_src = self.linear_src(node_features)  # (N, out_dim)
+        h_dst = self.linear_dst(node_features)  # (N, out_dim)
+
+        # Expand for pairwise computation
+        h_src_exp = jnp.expand_dims(h_src, axis=1)  # (N, 1, out_dim)
+        h_dst_exp = jnp.expand_dims(h_dst, axis=0)  # (1, N, out_dim)
+
+        # Broadcast to (N, N, out_dim)
+        h_src_broad = jnp.broadcast_to(h_src_exp, (num_nodes, num_nodes, h_src.shape[-1]))
+        h_dst_broad = jnp.broadcast_to(h_dst_exp, (num_nodes, num_nodes, h_dst.shape[-1]))
+
+        # Concatenate for attention
+        if self.use_edge_features and edge_features is not None:
+            h_edge = self.linear_edge(edge_features)  # (N, N, out_dim)
+            attn_input = jnp.concatenate([h_src_broad, h_dst_broad, h_edge], axis=-1)
+        else:
+            attn_input = jnp.concatenate([h_src_broad, h_dst_broad], axis=-1)
+
+        # Compute attention logits
+        attn_logits = self.attn(attn_input).squeeze(-1)  # (N, N)
+
+        # Apply LeakyReLU
+        attn_logits = jnp.where(
+            attn_logits >= 0,
+            attn_logits,
+            self.negative_slope * attn_logits,
+        )
+
+        # Mask non-edges with large negative value
+        attn_logits = jnp.where(adjacency > 0, attn_logits, -1e9)
+
+        # Softmax over neighbors
+        attn_weights = jax.nn.softmax(attn_logits, axis=-1)  # (N, N)
+
+        # Mask attention weights for non-edges
+        attn_weights = attn_weights * adjacency
+
+        # Aggregate: weighted sum of transformed features
+        out = jnp.einsum("ij,jd->id", attn_weights, h_dst)  # (N, out_dim)
+
+        return out, attn_weights
+
+
+class AttentiveFP(OperatorModule):
+    """AttentiveFP: Attention-based molecular fingerprint.
+
+    Implements the AttentiveFP architecture with:
+        1. Atom-level attention layers with GRU refinement
+        2. Molecule-level aggregation with attention and GRU
+        3. Final projection to fingerprint dimension
+
+    The model provides interpretable attention weights that indicate
+    which atoms contribute most to the molecular representation.
+
+    Example:
+        ```python
+        config = AttentiveFPConfig(hidden_dim=128, out_dim=256)
+        afp = AttentiveFP(config, rngs=nnx.Rngs(42))
+        data = {"node_features": nodes, "adjacency": adj, "edge_features": edges}
+        result, _, _ = afp.apply(data, {}, None)
+        fingerprint = result["fingerprint"]  # (256,)
+        attn = result["attention_weights"]  # interpretability
+        ```
+
+    References:
+        - Xiong et al. JCIM 2019
+    """
+
+    def __init__(
+        self,
+        config: AttentiveFPConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+    ):
+        """Initialize AttentiveFP.
+
+        Args:
+            config: AttentiveFP configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+
+        stabilize_operator_id(self)
+
+        if rngs is None:
+            rngs = nnx.Rngs(0)
+
+        # Initial linear projection
+        self.input_proj = nnx.Linear(config.in_features, config.hidden_dim, rngs=rngs)
+
+        # Atom-level: GATE convolutions with GRU
+        atom_convs = []
+        atom_grus = []
+
+        for _ in range(config.num_layers):
+            in_dim = config.hidden_dim
+            atom_convs.append(
+                GATEConv(
+                    in_dim=in_dim,
+                    out_dim=config.hidden_dim,
+                    edge_dim=config.edge_dim,
+                    negative_slope=config.negative_slope,
+                    rngs=rngs,
+                )
+            )
+            atom_grus.append(
+                nnx.GRUCell(
+                    in_features=config.hidden_dim,
+                    hidden_features=config.hidden_dim,
+                    rngs=rngs,
+                )
+            )
+        self.atom_convs = nnx.List(atom_convs)
+        self.atom_grus = nnx.List(atom_grus)
+
+        # Molecule-level aggregation
+        # Attention for global pooling
+        self.mol_attn = nnx.Linear(config.hidden_dim, 1, rngs=rngs)
+
+        # GRU for molecule-level refinement
+        mol_grus = []
+        for _ in range(config.num_timesteps):
+            mol_grus.append(
+                nnx.GRUCell(
+                    in_features=config.hidden_dim,
+                    hidden_features=config.hidden_dim,
+                    rngs=rngs,
+                )
+            )
+        self.mol_grus = nnx.List(mol_grus)
+
+        # Final projection
+        self.output_proj = nnx.Linear(config.hidden_dim, config.out_dim, rngs=rngs)
+
+        # Dropout
+        self.dropout = build_optional_dropout(config.dropout_rate, rngs=rngs)
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = 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.
+
+        Args:
+            data: Input data containing:
+                - node_features: (num_nodes, in_features) atom features
+                - adjacency: (num_nodes, num_nodes) adjacency matrix
+                - edge_features: Optional (num_nodes, num_nodes, edge_dim)
+                - node_mask: (num_nodes,) optional mask for valid nodes
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with "fingerprint" and "attention_weights" keys
+                - unchanged state
+                - unchanged metadata
+        """
+        del random_params, stats  # Unused
+
+        node_features = data["node_features"]
+        adjacency = data["adjacency"]
+        edge_features = data.get("edge_features")
+        node_mask = data.get("node_mask")
+
+        # Initial projection
+        h = self.input_proj(node_features)  # (N, hidden_dim)
+
+        if self.dropout is not None:
+            h = self.dropout(h)
+
+        # Collect attention weights for interpretability
+        all_attention_weights = []
+
+        # Atom-level message passing with GRU
+        for conv, gru in zip(self.atom_convs, self.atom_grus):
+            # Graph attention convolution
+            h_new, attn_weights = conv(h, adjacency, edge_features)
+            all_attention_weights.append(attn_weights)
+
+            if self.dropout is not None:
+                h_new = self.dropout(h_new)
+
+            # GRU update: h_new is input, h is hidden state
+            # GRUCell returns (new_carry, output) tuple - we use new_carry as next h
+            h, _ = gru(h, h_new)
+
+        # Apply node mask
+        if node_mask is not None:
+            h = h * node_mask[:, None]
+
+        # Molecule-level aggregation with attention
+        # Compute attention scores for global pooling
+        attn_scores = self.mol_attn(h).squeeze(-1)  # (N,)
+        if node_mask is not None:
+            attn_scores = jnp.where(node_mask > 0, attn_scores, -1e9)
+        attn_probs = jax.nn.softmax(attn_scores)  # (N,)
+
+        # Weighted sum for initial molecule representation
+        mol_repr = jnp.einsum("n,nd->d", attn_probs, h)  # (hidden_dim,)
+
+        # GRU refinement at molecule level
+        for mol_gru in self.mol_grus:
+            # Use atom representations as context
+            context = jnp.einsum("n,nd->d", attn_probs, h)
+            # GRUCell returns (new_carry, output) tuple
+            mol_repr, _ = mol_gru(mol_repr, context)
+
+        # Final projection
+        fingerprint = self.output_proj(mol_repr)
+
+        if self.dropout is not None:
+            fingerprint = self.dropout(fingerprint)
+
+        result = {
+            **data,
+            "fingerprint": fingerprint,
+            "attention_weights": all_attention_weights,
+            "molecule_attention": attn_probs,
+        }
+
+        return result, state, metadata
+
+
+def create_attentive_fp(
+    hidden_dim: int = 200,
+    out_dim: int = 200,
+    num_layers: int = 2,
+    num_timesteps: int = 2,
+    dropout_rate: float = 0.0,
+    seed: int = 42,
+) -> AttentiveFP:
+    """Create an AttentiveFP operator.
+
+    Args:
+        hidden_dim: Hidden dimension for GNN layers.
+        out_dim: Output fingerprint dimension.
+        num_layers: Number of atom-level attention layers.
+        num_timesteps: Number of molecule-level GRU iterations.
+        dropout_rate: Dropout rate.
+        seed: Random seed.
+
+    Returns:
+        Configured AttentiveFP.
+    """
+    config = AttentiveFPConfig(
+        hidden_dim=hidden_dim,
+        out_dim=out_dim,
+        num_layers=num_layers,
+        num_timesteps=num_timesteps,
+        dropout_rate=dropout_rate,
+    )
+    return AttentiveFP(config, rngs=nnx.Rngs(seed))