torchcell.models.diffusion_decoder.CrossAttention

class CrossAttention(dim: int, num_heads: int = 8, dropout: float = 0.1)

Bases: Module

Cross-attention module where diffusion state queries graph embeddings.

forward(x: Tensor, context: Tensor) → Tensor

Attend the diffusion state to the graph-embedding context.