quantmllibrary 0.1.0__py3-none-any.whl

quantml/models/__init__.py

@@ -0,0 +1,27 @@
+"""
+QuantML Models
+
+This module provides neural network models optimized for quantitative trading.
+"""
+
+from quantml.models.linear import Linear
+from quantml.models.rnn import SimpleRNN
+from quantml.models.tcn import TCN, TCNBlock
+from quantml.models.lstm import LSTM, LSTMCell
+from quantml.models.gru import GRU, GRUCell
+from quantml.models.mlp import MLP, ResidualMLP, create_mlp
+from quantml.models.normalization import BatchNorm1d, LayerNorm
+from quantml.models.dropout import Dropout
+from quantml.models.attention import SelfAttention, MultiHeadAttention
+
+__all__ = [
+    'Linear', 
+    'SimpleRNN', 
+    'TCN', 'TCNBlock',
+    'LSTM', 'LSTMCell',
+    'GRU', 'GRUCell',
+    'MLP', 'ResidualMLP', 'create_mlp',
+    'BatchNorm1d', 'LayerNorm',
+    'Dropout',
+    'SelfAttention', 'MultiHeadAttention',
+]

quantml/models/attention.py

@@ -0,0 +1,258 @@
+"""
+Attention mechanisms.
+
+Implementations of Self-Attention and Multi-Head Attention.
+"""
+
+from typing import Optional, List, Tuple
+import math
+from quantml.tensor import Tensor
+from quantml import ops
+from quantml.models.linear import Linear
+
+
+class SelfAttention:
+    """
+    Scaled Dot-Product Self-Attention.
+    
+    Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
+    
+    Attributes:
+        embed_dim: Dimension of embeddings
+        num_heads: Number of attention heads (default: 1)
+        dropout: Dropout probability
+    
+    Examples:
+        >>> attn = SelfAttention(64)
+        >>> x = Tensor([[[1.0] * 64] * 10])  # batch x seq x dim
+        >>> out = attn.forward(x)
+    """
+    
+    def __init__(self, embed_dim: int, dropout: float = 0.0):
+        self.embed_dim = embed_dim
+        self.dropout = dropout
+        
+        # Projections
+        self.q_proj = Linear(embed_dim, embed_dim)
+        self.k_proj = Linear(embed_dim, embed_dim)
+        self.v_proj = Linear(embed_dim, embed_dim)
+        self.out_proj = Linear(embed_dim, embed_dim)
+        
+        # Scaling factor
+        self.scale = 1.0 / math.sqrt(embed_dim)
+        
+    def forward(
+        self, 
+        x: Tensor, 
+        mask: Optional[Tensor] = None
+    ) -> Tensor:
+        """
+        Forward pass.
+        
+        Args:
+            x: Input tensor (batch x seq_len x embed_dim)
+            mask: Optional mask (batch x seq_len x seq_len)
+        
+        Returns:
+            Output tensor (batch x seq_len x embed_dim)
+        """
+        # Linear projections
+        # Each is (batch x seq_len x embed_dim)
+        q = self.q_proj.forward(x)
+        k = self.k_proj.forward(x)
+        v = self.v_proj.forward(x)
+        
+        # Calculate attention scores: Q @ K^T
+        # Transpose k for matmul
+        # Q: (B, S, D), K: (B, S, D) -> K^T: (B, D, S)
+        # We need an explicit transpose op for 3D tensors which we don't have yet
+        # So we'll implement a simplified version assuming batch=1 or flattened
+        
+        # NOTE: Full multi-head attention with batching requires reshaping and 3D matmul
+        # which depends on robust tensor ops. This is a simplified implementation.
+        
+        # Fallback to pure python loop for batch support
+        data = x.data
+        if isinstance(data[0][0], list): # 3D
+            batch_size = len(data)
+            out_batches = []
+            
+            for b in range(batch_size):
+                # Extract batch slice as 2D tensor
+                q_b = self._get_batch_slice(q, b)
+                k_b = self._get_batch_slice(k, b)
+                v_b = self._get_batch_slice(v, b)
+                mask_b = self._get_batch_slice(mask, b) if mask is not None else None
+                
+                out_b = self._attention_2d(q_b, k_b, v_b, mask_b)
+                out_batches.append(out_b.data)
+            
+            # Combine back to 3D tensor
+            # Note: This breaks autograd graph effectively unless we are careful
+            # Ideally we'd use 3D ops
+            attn_out = Tensor(out_batches, requires_grad=True) # Graph broken here for now
+        else:
+            # 2D case
+            attn_out = self._attention_2d(q, k, v, mask)
+            
+        # Output projection
+        return self.out_proj.forward(attn_out)
+        
+    def _attention_2d(self, An: Tensor, K: Tensor, V: Tensor, mask: Optional[Tensor]) -> Tensor:
+        """Compute attention for single sample (seq_len x dim)."""
+        # K^T
+        # K is (S, D), we need (D, S) manually transposed
+        K_T = self._transpose_2d(K)
+        
+        # Scores: (S, D) @ (D, S) -> (S, S)
+        scores = ops.matmul(An, K_T)
+        
+        # Scale
+        scaled_scores = ops.mul(scores, self.scale)
+        
+        # Mask
+        if mask is not None:
+            # Assume mask is additive (0 for keep, -inf for mask)
+            scaled_scores = ops.add(scaled_scores, mask)
+            
+        # Softmax over last dim (rows)
+        attn_weights = ops.softmax(scaled_scores, axis=-1)
+        
+        # Output: (S, S) @ (S, D) -> (S, D)
+        output = ops.matmul(attn_weights, V)
+        
+        return output
+
+    def _transpose_2d(self, t: Tensor) -> Tensor:
+        """Transpose 2D tensor."""
+        data = t.data
+        rows = len(data)
+        cols = len(data[0])
+        new_data = [[data[i][j] for i in range(rows)] for j in range(cols)]
+        return Tensor(new_data, requires_grad=t.requires_grad)
+
+    def _get_batch_slice(self, t: Tensor, idx: int) -> Tensor:
+        """Get 2D slice from 3D tensor."""
+        return Tensor(t.data[idx], requires_grad=t.requires_grad)
+        
+    def parameters(self) -> List[Tensor]:
+        """Get parameters."""
+        return (self.q_proj.parameters() + 
+                self.k_proj.parameters() + 
+                self.v_proj.parameters() + 
+                self.out_proj.parameters())
+                
+    def zero_grad(self) -> None:
+        """Zero gradients."""
+        for p in self.parameters():
+            p.zero_grad()
+
+
+class MultiHeadAttention:
+    """
+    Multi-Head Attention.
+    
+    Splits embedding into multiple heads, applies attention independently,
+    and concatenates results.
+    
+    Attributes:
+        embed_dim: Model dimension
+        num_heads: Number of heads
+    """
+    
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
+        if embed_dim % num_heads != 0:
+            raise ValueError(f"embed_dim {embed_dim} must be divisible by num_heads {num_heads}")
+            
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        
+        # In typical implementation, we project to (num_heads * head_dim)
+        # which is same as embed_dim. So we can use one big Linear layer
+        # and reshape/split logic.
+        
+        self.q_proj = Linear(embed_dim, embed_dim)
+        self.k_proj = Linear(embed_dim, embed_dim)
+        self.v_proj = Linear(embed_dim, embed_dim)
+        self.out_proj = Linear(embed_dim, embed_dim)
+        
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        
+    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        """
+        Simple implementation: Loop over heads (inefficient but clear).
+        True parallel implementation requires robust reshape/transpose ops.
+        """
+        batch_size = len(x.data) if isinstance(x.data[0], list) and isinstance(x.data[0][0], list) else 1
+        
+        # 1. Projections
+        q = self.q_proj.forward(x)
+        k = self.k_proj.forward(x)
+        v = self.v_proj.forward(x)
+        
+        # 2. Split heads and generic attention logic is complex without proper tensor reshaping
+        # For this library's current state, we will implement a simplified single-head equivalent
+        # that mathematically matches but doesn't actually split distinct subspaces without reshaping.
+        
+        # For now, we'll delegate to a simpler attention mechanism that treats it as one big head
+        # This is a PLACEHOLDER until complex tensor manipulation functions (view, permute) are added.
+        
+        # We'll re-use the SelfAttention logic which effectively does 1 head
+        # To support multi-head properly, we need ops.reshape / ops.transpose(permute)
+        
+        return self._fallback_single_head(q, k, v, mask)
+
+    def _fallback_single_head(self, q, k, v, mask):
+        # Same logic as SelfAttention._attention_2d but handling batches loop
+        data = q.data
+        if isinstance(data[0][0], list): # 3D
+            batch_size = len(data)
+            out_batches = []
+            for b in range(batch_size):
+                q_b = self._get_batch_slice(q, b)
+                k_b = self._get_batch_slice(k, b)
+                v_b = self._get_batch_slice(v, b)
+                mask_b = self._get_batch_slice(mask, b) if mask is not None else None
+                
+                # Manual attention with correct scaling for head_dim
+                K_T = self._transpose_2d(k_b)
+                scores = ops.matmul(q_b, K_T)
+                scaled_scores = ops.mul(scores, self.scale)
+                if mask_b is not None:
+                    scaled_scores = ops.add(scaled_scores, mask_b)
+                attn_weights = ops.softmax(scaled_scores, axis=-1)
+                out_b = ops.matmul(attn_weights, v_b)
+                
+                out_batches.append(out_b.data)
+            attn_out = Tensor(out_batches, requires_grad=True)
+        else:
+            K_T = self._transpose_2d(k)
+            scores = ops.matmul(q, K_T)
+            scaled_scores = ops.mul(scores, self.scale)
+            if mask is not None:
+                scaled_scores = ops.add(scaled_scores, mask)
+            attn_weights = ops.softmax(scaled_scores, axis=-1)
+            attn_out = ops.matmul(attn_weights, v)
+            
+        return self.out_proj.forward(attn_out)
+
+    def _transpose_2d(self, t: Tensor) -> Tensor:
+        data = t.data
+        rows = len(data)
+        cols = len(data[0])
+        new_data = [[data[i][j] for i in range(rows)] for j in range(cols)]
+        return Tensor(new_data, requires_grad=t.requires_grad)
+
+    def _get_batch_slice(self, t: Tensor, idx: int) -> Tensor:
+        return Tensor(t.data[idx], requires_grad=t.requires_grad)
+
+    def parameters(self) -> List[Tensor]:
+        return (self.q_proj.parameters() + 
+                self.k_proj.parameters() + 
+                self.v_proj.parameters() + 
+                self.out_proj.parameters())
+    
+    def zero_grad(self) -> None:
+        for p in self.parameters():
+            p.zero_grad()

quantml/models/dropout.py

@@ -0,0 +1,130 @@
+"""
+Dropout regularization.
+
+Implements Dropout layer for preventing overfitting.
+"""
+
+from typing import Optional, List
+import random
+from quantml.tensor import Tensor
+from quantml import ops
+
+
+class Dropout:
+    """
+    Dropout layer.
+    
+    Randomly zeroes some elements of the input tensor with probability p
+    using samples from a Bernoulli distribution.
+    
+    During training, outputs are scaled by 1/(1-p).
+    During evaluation, does nothing.
+    
+    Attributes:
+        p: Probability of an element being zeroed
+        inplace: If True, do operation in-place (not supported yet)
+    
+    Examples:
+        >>> dropout = Dropout(p=0.5)
+        >>> x = Tensor([[1.0, 2.0, 3.0]])
+        >>> y = dropout(x)
+    """
+    
+    def __init__(self, p: float = 0.5, inplace: bool = False):
+        if p < 0 or p > 1:
+            raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+    
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Forward pass.
+        
+        Args:
+            x: Input tensor
+        
+        Returns:
+            Tensor with dropout applied
+        """
+        if not self.training or self.p == 0:
+            return x
+        
+        # Apply dropout mask
+        scale = 1.0 / (1.0 - self.p)
+        
+        data = x.data
+        if isinstance(data[0], list):
+            # 2D case (common for dropout)
+            if isinstance(data[0][0], list):
+                # 3D
+                out_data = []
+                for i in range(len(data)):
+                    batch_out = []
+                    for j in range(len(data[i])):
+                        row_out = []
+                        for k in range(len(data[i][j])):
+                            if random.random() > self.p:
+                                row_out.append(float(data[i][j][k]) * scale)
+                            else:
+                                row_out.append(0.0)
+                        batch_out.append(row_out)
+                    out_data.append(batch_out)
+            else:
+                # 2D
+                out_data = []
+                for i in range(len(data)):
+                    row_out = []
+                    for j in range(len(data[i])):
+                        if random.random() > self.p:
+                            row_out.append(float(data[i][j]) * scale)
+                        else:
+                            row_out.append(0.0)
+                    out_data.append(row_out)
+        else:
+            # 1D
+            out_data = []
+            for i in range(len(data)):
+                if random.random() > self.p:
+                    out_data.append(float(data[i]) * scale)
+                else:
+                    out_data.append(0.0)
+        
+        # Note: Dropout backward is handled because we are returning a new Tensor
+        # However, for correct autograd, we should probably implement it as an op
+        # Or multiply by a mask tensor
+        
+        # Better approach for autograd: Create a mask tensor and multiply
+        # This ensures the graph captures the dropout operation correctly
+        
+        mask_data = self._generate_mask_data(x.data, scale)
+        mask = Tensor(mask_data, requires_grad=False)
+        
+        return ops.mul(x, mask)
+    
+    def _generate_mask_data(self, data, scale):
+        """Recursively generate dropout mask."""
+        if isinstance(data, list):
+            return [self._generate_mask_data(item, scale) for item in data]
+        else:
+            return scale if random.random() > self.p else 0.0
+    
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.forward(x)
+    
+    def train(self, mode: bool = True) -> 'Dropout':
+        """Set training mode."""
+        self.training = mode
+        return self
+    
+    def eval(self) -> 'Dropout':
+        """Set evaluation mode."""
+        return self.train(False)
+    
+    def parameters(self) -> List[Tensor]:
+        """Get trainable parameters (none for dropout)."""
+        return []
+        
+    def zero_grad(self) -> None:
+        """Clear gradients (no-op)."""
+        pass