rustytorch/crates/training/rtx-transformers/src/graph/mod.rs

//! Graph Transformer components for graph-structured data
//!
//! This module provides attention-based learning capabilities on graph-structured data,
//! supporting both homogeneous and heterogeneous graphs with batch processing.

use crate::Result;
use rtx_tensor::{Device, Tensor};

pub mod graph_attention_simple;
pub mod graph_pooling_simple;
pub mod graph_transformer_simple;
pub mod positional_encoding_simple;

#[cfg(all(test, feature = "disabled_tests"))]
pub mod graph_transformer_tests;

pub use graph_attention_simple::{AttentionOutput, GraphAttention, GraphAttentionConfig};
pub use graph_pooling_simple::{
    AttentionPooling, GlobalPooling, GraphPooling, HierarchicalPooling, Set2SetPooling,
};
pub use graph_transformer_simple::{GraphOutput, GraphTransformer, GraphTransformerConfig};
pub use positional_encoding_simple::{GraphPEConfig, GraphPEType, GraphPositionalEncoding};

/// Graph data structure for batch processing
#[derive(Debug, Clone)]
pub struct GraphBatch {
    /// Node features: [total_nodes, node_dim]
    pub node_features: Tensor,
    /// Edge features: [total_edges, edge_dim]
    pub edge_features: Tensor,
    /// Edge indices: [2, total_edges] (source, target)
    pub edge_indices: Tensor,
    /// Graph boundaries for batch processing: [batch_size + 1]
    pub graph_boundaries: Vec<usize>,
    /// Number of graphs in batch
    pub batch_size: usize,
    /// Device
    pub device: Device,
}

impl GraphBatch {
    /// Create a new graph batch
    pub fn new(
        node_features: Tensor,
        edge_features: Tensor,
        edge_indices: Tensor,
        graph_boundaries: Vec<usize>,
        device: Device,
    ) -> Result<Self> {
        let batch_size = graph_boundaries.len().saturating_sub(1);

        Ok(Self {
            node_features,
            edge_features,
            edge_indices,
            graph_boundaries,
            batch_size,
            device,
        })
    }

    /// Get the number of nodes in the batch
    pub fn num_nodes(&self) -> usize {
        self.node_features.shape().dims()[0]
    }

    /// Get the number of edges in the batch
    pub fn num_edges(&self) -> usize {
        self.edge_features.shape().dims()[0]
    }

    /// Check if batch is homogeneous (all graphs have same structure)
    pub fn is_homogeneous(&self) -> bool {
        // For now, assume heterogeneous - can be optimized later
        false
    }
}

/// Graph attention mechanism types
#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum GraphAttentionType {
    /// Standard graph attention (GAT)
    Standard,
    /// Graph attention with edge features
    EdgeAware,
    /// Multi-head graph attention
    MultiHead,
    /// Graph attention with gating mechanism
    Gated,
}

/// Attention mechanism types
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum AttentionMechanism {
    /// Dot-product attention
    DotProduct,
    /// Additive attention
    Additive,
    /// Scaled dot-product attention
    ScaledDotProduct,
}

/// Graph pooling strategies
#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum PoolingStrategy {
    /// Global mean/max/sum pooling
    Global,
    /// Hierarchical graph pooling
    Hierarchical,
    /// Set2Set pooling mechanism
    Set2Set,
    /// Attention-based pooling
    Attention,
}

/// Graph transformer layer trait
pub trait GraphLayer: Send + Sync {
    /// Forward pass through the graph layer
    fn forward(&self, graph: &GraphBatch) -> Result<GraphOutput>;

    /// Get the layer type name
    fn layer_type(&self) -> &'static str;

    /// Get the device this layer is on
    fn device(&self) -> &Device;

    /// Get layer parameters for optimization
    fn parameters(&self) -> Vec<&Tensor>;

    /// Get mutable layer parameters for optimization
    fn parameters_mut(&mut self) -> Vec<&mut Tensor>;
}