openarchx 0.1.0__py3-none-any.whl

openarchx/nn/transformers.py

@@ -0,0 +1,187 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+from .layers import Linear, LayerNorm
+from .activations import ReLU
+
+class MultiheadAttention(Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+        
+        self.q_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self.k_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+
+    def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
+        batch_size = query.data.shape[0]
+        
+        # Linear projections and reshape
+        q = self._reshape_for_attention(self.q_proj(query))
+        k = self._reshape_for_attention(self.k_proj(key))
+        v = self._reshape_for_attention(self.v_proj(value))
+        
+        # Scaled dot-product attention
+        scaling = float(self.head_dim) ** -0.5
+        attn = np.matmul(q.data, k.data.transpose(0, 1, 3, 2)) * scaling
+        
+        if attn_mask is not None:
+            attn = np.where(attn_mask, -np.inf, attn)
+        
+        attn = self._softmax(attn)
+        
+        if self.dropout > 0:
+            attn = np.where(np.random.random(attn.shape) > self.dropout, attn, 0) / (1 - self.dropout)
+        
+        output = np.matmul(attn, v.data)
+        output = output.transpose(0, 2, 1, 3).reshape(batch_size, -1, self.embed_dim)
+        
+        return self.out_proj(Tensor(output, requires_grad=True))
+
+    def _reshape_for_attention(self, x):
+        batch_size, seq_len, _ = x.data.shape
+        return Tensor(
+            x.data.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+            .transpose(0, 2, 1, 3),
+            requires_grad=True
+        )
+
+    def _softmax(self, x):
+        exp_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
+        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
+
+class TransformerEncoderLayer(Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.linear2 = Linear(dim_feedforward, d_model)
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout = dropout
+        self.activation = ReLU()
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        # Self attention block
+        src2 = self.self_attn(src, src, src, attn_mask=src_mask,
+                             key_padding_mask=src_key_padding_mask)
+        src = src + self._dropout_layer(src2)
+        src = self.norm1(src)
+        
+        # Feedforward block
+        src2 = self.linear2(self._dropout_layer(self.activation(self.linear1(src))))
+        src = src + self._dropout_layer(src2)
+        src = self.norm2(src)
+        
+        return src
+
+    def _dropout_layer(self, x):
+        if self.dropout > 0:
+            mask = np.random.random(x.data.shape) > self.dropout
+            return Tensor(mask * x.data / (1 - self.dropout), requires_grad=True)
+        return x
+
+class TransformerDecoderLayer(Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
+        super().__init__()
+        self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.linear2 = Linear(dim_feedforward, d_model)
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.norm3 = LayerNorm(d_model)
+        self.dropout = dropout
+        self.activation = ReLU()
+
+    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        # Self attention block
+        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
+                             key_padding_mask=tgt_key_padding_mask)
+        tgt = tgt + self._dropout_layer(tgt2)
+        tgt = self.norm1(tgt)
+        
+        # Cross attention block
+        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
+                                  key_padding_mask=memory_key_padding_mask)
+        tgt = tgt + self._dropout_layer(tgt2)
+        tgt = self.norm2(tgt)
+        
+        # Feedforward block
+        tgt2 = self.linear2(self._dropout_layer(self.activation(self.linear1(tgt))))
+        tgt = tgt + self._dropout_layer(tgt2)
+        tgt = self.norm3(tgt)
+        
+        return tgt
+
+    def _dropout_layer(self, x):
+        if self.dropout > 0:
+            mask = np.random.random(x.data.shape) > self.dropout
+            return Tensor(mask * x.data / (1 - self.dropout), requires_grad=True)
+        return x
+
+class TransformerEncoder(Module):
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = [encoder_layer for _ in range(num_layers)]
+        self.norm = norm
+
+    def forward(self, src, mask=None, src_key_padding_mask=None):
+        output = src
+        for layer in self.layers:
+            output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)
+        
+        if self.norm is not None:
+            output = self.norm(output)
+        
+        return output
+
+class TransformerDecoder(Module):
+    def __init__(self, decoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = [decoder_layer for _ in range(num_layers)]
+        self.norm = norm
+
+    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
+                tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        output = tgt
+        for layer in self.layers:
+            output = layer(output, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
+                         tgt_key_padding_mask=tgt_key_padding_mask,
+                         memory_key_padding_mask=memory_key_padding_mask)
+        
+        if self.norm is not None:
+            output = self.norm(output)
+        
+        return output
+
+class Transformer(Module):
+    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1):
+        super().__init__()
+        
+        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
+        encoder_norm = LayerNorm(d_model)
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+        
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
+        decoder_norm = LayerNorm(d_model)
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
+        
+        self.d_model = d_model
+        self.nhead = nhead
+
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None,
+                src_key_padding_mask=None, tgt_key_padding_mask=None,
+                memory_key_padding_mask=None):
+        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
+        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
+                            tgt_key_padding_mask=tgt_key_padding_mask,
+                            memory_key_padding_mask=memory_key_padding_mask)
+        return output

openarchx/optimizers/adam.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+class Adam:
+    def __init__(self, parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, clip_grad=None):
+        self.parameters = parameters
+        self.lr = lr
+        self.beta1, self.beta2 = betas
+        self.eps = eps
+        self.weight_decay = weight_decay
+        self.clip_grad = clip_grad
+        
+        # Initialize momentum and velocity
+        self.m = [np.zeros_like(param.data) for param in parameters]
+        self.v = [np.zeros_like(param.data) for param in parameters]
+        self.t = 0
+
+    def clip_gradients(self):
+        if self.clip_grad is not None:
+            for param in self.parameters:
+                if param.grad is not None:
+                    np.clip(param.grad, -self.clip_grad, self.clip_grad, out=param.grad)
+
+    def step(self):
+        self.t += 1
+        self.clip_gradients()
+        
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                # Update momentum
+                self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
+                
+                # Update velocity
+                self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * (grad ** 2)
+                
+                # Bias correction
+                m_hat = self.m[i] / (1 - self.beta1 ** self.t)
+                v_hat = self.v[i] / (1 - self.beta2 ** self.t)
+                
+                # Update parameters
+                param.data -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
+    
+    def zero_grad(self):
+        for param in self.parameters:
+            if param.grad is not None:
+                param.grad.fill(0)

openarchx/optimizers/adaptive.py

@@ -0,0 +1,63 @@
+import numpy as np
+from .base import Optimizer
+
+class Adagrad(Optimizer):
+    """Adaptive Gradient Algorithm"""
+    def __init__(self, parameters, lr=0.01, eps=1e-8, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.eps = eps
+        self.G = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                self.G[i] += np.square(grad)
+                param.data -= self.lr * grad / (np.sqrt(self.G[i]) + self.eps)
+
+class Adadelta(Optimizer):
+    """Adaptive Delta Algorithm"""
+    def __init__(self, parameters, rho=0.95, eps=1e-6, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, 1.0, weight_decay, clip_grad)  # lr=1.0 as it's not used
+        self.rho = rho
+        self.eps = eps
+        self.G = [np.zeros_like(param.data) for param in parameters]
+        self.delta = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                self.G[i] = self.rho * self.G[i] + (1 - self.rho) * np.square(grad)
+                rms_g = np.sqrt(self.G[i] + self.eps)
+                rms_delta = np.sqrt(self.delta[i] + self.eps)
+                
+                update = -rms_delta / rms_g * grad
+                param.data += update
+                
+                self.delta[i] = self.rho * self.delta[i] + (1 - self.rho) * np.square(update)
+
+class RMSprop(Optimizer):
+    """Root Mean Square Propagation"""
+    def __init__(self, parameters, lr=0.01, alpha=0.99, eps=1e-8, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.alpha = alpha
+        self.eps = eps
+        self.G = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                self.G[i] = self.alpha * self.G[i] + (1 - self.alpha) * np.square(grad)
+                param.data -= self.lr * grad / (np.sqrt(self.G[i]) + self.eps)

openarchx/optimizers/base.py

@@ -0,0 +1,24 @@
+import numpy as np
+from abc import ABC, abstractmethod
+
+class Optimizer(ABC):
+    def __init__(self, parameters, lr=0.001, weight_decay=0.0, clip_grad=None):
+        self.parameters = parameters
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.clip_grad = clip_grad
+
+    def clip_gradients(self):
+        if self.clip_grad is not None:
+            for param in self.parameters:
+                if param.grad is not None:
+                    np.clip(param.grad, -self.clip_grad, self.clip_grad, out=param.grad)
+
+    @abstractmethod
+    def step(self):
+        pass
+
+    def zero_grad(self):
+        for param in self.parameters:
+            if param.grad is not None:
+                param.grad.fill(0)

openarchx/optimizers/modern.py

@@ -0,0 +1,98 @@
+import numpy as np
+from .base import Optimizer
+
+class RAdam(Optimizer):
+    """Rectified Adam Optimizer"""
+    def __init__(self, parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.betas = betas
+        self.eps = eps
+        self.m = [np.zeros_like(param.data) for param in parameters]
+        self.v = [np.zeros_like(param.data) for param in parameters]
+        self.step_count = 0
+        
+    def step(self):
+        self.step_count += 1
+        self.clip_gradients()
+        
+        beta1, beta2 = self.betas
+        bias_correction1 = 1 - beta1 ** self.step_count
+        bias_correction2 = 1 - beta2 ** self.step_count
+        
+        rho_inf = 2 / (1 - beta2) - 1
+        rho_t = rho_inf - 2 * self.step_count * beta2 ** self.step_count / bias_correction2
+        
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                self.m[i] = beta1 * self.m[i] + (1 - beta1) * grad
+                self.v[i] = beta2 * self.v[i] + (1 - beta2) * np.square(grad)
+                
+                m_hat = self.m[i] / bias_correction1
+                
+                if rho_t > 4:
+                    r = np.sqrt(((rho_t - 4) * (rho_t - 2) * rho_inf) / ((rho_inf - 4) * (rho_inf - 2) * rho_t))
+                    v_hat = np.sqrt(self.v[i] / bias_correction2)
+                    param.data -= self.lr * r * m_hat / (v_hat + self.eps)
+                else:
+                    param.data -= self.lr * m_hat
+
+class AdaBelief(Optimizer):
+    """AdaBelief Optimizer"""
+    def __init__(self, parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-16, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.betas = betas
+        self.eps = eps
+        self.m = [np.zeros_like(param.data) for param in parameters]
+        self.s = [np.zeros_like(param.data) for param in parameters]
+        self.step_count = 0
+
+    def step(self):
+        self.step_count += 1
+        self.clip_gradients()
+        
+        beta1, beta2 = self.betas
+        bias_correction1 = 1 - beta1 ** self.step_count
+        bias_correction2 = 1 - beta2 ** self.step_count
+        
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                self.m[i] = beta1 * self.m[i] + (1 - beta1) * grad
+                diff = grad - self.m[i]
+                self.s[i] = beta2 * self.s[i] + (1 - beta2) * np.square(diff)
+                
+                m_hat = self.m[i] / bias_correction1
+                s_hat = self.s[i] / bias_correction2
+                
+                param.data -= self.lr * m_hat / (np.sqrt(s_hat) + self.eps)
+
+class Lion(Optimizer):
+    """Lion Optimizer (Learning with Inner Optimization)"""
+    def __init__(self, parameters, lr=0.0001, betas=(0.9, 0.99), weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.betas = betas
+        self.m = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        beta1, beta2 = self.betas
+        
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                update = beta1 * self.m[i] + (1 - beta1) * grad
+                old_m = self.m[i].copy()
+                self.m[i] = beta2 * self.m[i] + (1 - beta2) * grad
+                
+                # Update using sign of momentum
+                param.data -= self.lr * np.sign(update)

openarchx/optimizers/optx.py

@@ -0,0 +1,91 @@
+import numpy as np
+
+class OptX:
+    def __init__(self, parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0.0, clip_grad=None, rectify=True,
+                 lookahead_steps=5, lookahead_alpha=0.5):
+        self.parameters = parameters
+        self.lr = lr
+        self.beta1, self.beta2 = betas
+        self.eps = eps
+        self.weight_decay = weight_decay
+        self.clip_grad = clip_grad
+        self.rectify = rectify
+        self.lookahead_steps = lookahead_steps
+        self.lookahead_alpha = lookahead_alpha
+        
+        # Initialize momentum and velocity
+        self.m = [np.zeros_like(param.data) for param in parameters]
+        self.v = [np.zeros_like(param.data) for param in parameters]
+        self.slow_weights = [param.data.copy() for param in parameters]
+        self.t = 0
+        self.step_counter = 0
+        
+        # Gradient variance tracking
+        self.grad_var = [np.zeros_like(param.data) for param in parameters]
+        
+    def clip_gradients(self):
+        if self.clip_grad is not None:
+            for param in self.parameters:
+                if param.grad is not None:
+                    np.clip(param.grad, -self.clip_grad, self.clip_grad, out=param.grad)
+                    
+    def update_grad_variance(self, i, grad):
+        # Update running variance of gradients
+        if self.t > 1:
+            self.grad_var[i] = 0.9 * self.grad_var[i] + 0.1 * (grad - self.m[i])**2
+
+    def compute_adaptive_lr(self, i, v_hat):
+        # Compute adaptive learning rate based on gradient variance
+        if self.t > 1:
+            variance_scaling = 1.0 / (1.0 + np.sqrt(self.grad_var[i]) + self.eps)
+            return self.lr * variance_scaling
+        return self.lr
+
+    def step(self):
+        self.t += 1
+        self.step_counter += 1
+        self.clip_gradients()
+        
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                
+                # Update momentum with bias correction
+                self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
+                m_hat = self.m[i] / (1 - self.beta1 ** self.t)
+                
+                # Update velocity with bias correction
+                self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * (grad ** 2)
+                v_hat = self.v[i] / (1 - self.beta2 ** self.t)
+                
+                # Update gradient variance for adaptive scaling
+                self.update_grad_variance(i, grad)
+                
+                # Compute adaptive learning rate
+                adaptive_lr = self.compute_adaptive_lr(i, v_hat)
+                
+                # Compute update
+                update = adaptive_lr * m_hat / (np.sqrt(v_hat) + self.eps)
+                
+                if self.rectify:
+                    # Apply rectification to prevent overshooting
+                    variance_ratio = np.sqrt(self.grad_var[i]) / (np.sqrt(v_hat) + self.eps)
+                    update *= np.minimum(1.0, np.maximum(0.1, variance_ratio))
+                
+                # Apply update
+                param.data -= update
+                
+                # Lookahead update
+                if self.step_counter % self.lookahead_steps == 0:
+                    # Move slow weights toward current parameters
+                    self.slow_weights[i] += self.lookahead_alpha * (param.data - self.slow_weights[i])
+                    # Update parameters to interpolated position
+                    param.data = self.slow_weights[i].copy()
+    
+    def zero_grad(self):
+        for param in self.parameters:
+            if param.grad is not None:
+                param.grad.fill(0)

openarchx/optimizers/sgd.py

@@ -0,0 +1,63 @@
+import numpy as np
+from .base import Optimizer
+
+class BGD(Optimizer):
+    """Batch Gradient Descent"""
+    def step(self):
+        self.clip_gradients()
+        for param in self.parameters:
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                param.data -= self.lr * grad
+
+class MBGD(Optimizer):
+    """Mini-Batch Gradient Descent"""
+    def __init__(self, parameters, lr=0.001, batch_size=32, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.batch_size = batch_size
+
+    def step(self):
+        self.clip_gradients()
+        for param in self.parameters:
+            if param.grad is not None:
+                grad = param.grad / self.batch_size
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                param.data -= self.lr * grad
+
+class SGD(Optimizer):
+    """SGD with Momentum"""
+    def __init__(self, parameters, lr=0.001, momentum=0.9, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.momentum = momentum
+        self.v = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                self.v[i] = self.momentum * self.v[i] - self.lr * grad
+                param.data += self.v[i]
+
+class NAG(Optimizer):
+    """Nesterov Accelerated Gradient"""
+    def __init__(self, parameters, lr=0.001, momentum=0.9, weight_decay=0.0, clip_grad=None):
+        super().__init__(parameters, lr, weight_decay, clip_grad)
+        self.momentum = momentum
+        self.v = [np.zeros_like(param.data) for param in parameters]
+
+    def step(self):
+        self.clip_gradients()
+        for i, param in enumerate(self.parameters):
+            if param.grad is not None:
+                grad = param.grad
+                if self.weight_decay > 0:
+                    grad = grad + self.weight_decay * param.data
+                v_prev = self.v[i].copy()
+                self.v[i] = self.momentum * self.v[i] - self.lr * grad
+                param.data += -self.momentum * v_prev + (1 + self.momentum) * self.v[i]

openarchx/quantum/circuit.py

@@ -0,0 +1,92 @@
+import numpy as np
+from ..core.tensor import Tensor
+from ..nn.module import Module
+from .gates import QuantumGates, QuantumRegister
+
+class QuantumCircuit(Module):
+    def __init__(self, num_qubits):
+        super().__init__()
+        self.num_qubits = num_qubits
+        self.register = None
+        self.gates = []
+    
+    def reset(self):
+        """Initialize a new quantum register"""
+        self.register = QuantumRegister(self.num_qubits)
+        self.gates = []
+    
+    def h(self, qubit):
+        """Apply Hadamard gate"""
+        self.gates.append(('H', qubit))
+        return self
+    
+    def x(self, qubit):
+        """Apply Pauli-X gate"""
+        self.gates.append(('X', qubit))
+        return self
+    
+    def y(self, qubit):
+        """Apply Pauli-Y gate"""
+        self.gates.append(('Y', qubit))
+        return self
+    
+    def z(self, qubit):
+        """Apply Pauli-Z gate"""
+        self.gates.append(('Z', qubit))
+        return self
+    
+    def rx(self, theta, qubit):
+        """Apply parameterized rotation around X axis"""
+        self.gates.append(('RX', theta, qubit))
+        return self
+    
+    def ry(self, theta, qubit):
+        """Apply parameterized rotation around Y axis"""
+        self.gates.append(('RY', theta, qubit))
+        return self
+    
+    def rz(self, theta, qubit):
+        """Apply parameterized rotation around Z axis"""
+        self.gates.append(('RZ', theta, qubit))
+        return self
+    
+    def cnot(self, control, target):
+        """Apply CNOT gate"""
+        self.gates.append(('CNOT', control, target))
+        return self
+    
+    def forward(self, input_params=None):
+        """Execute quantum circuit and return measurement results"""
+        self.reset()
+        
+        # Apply gates with batch support
+        for gate in self.gates:
+            if gate[0] in ['RX', 'RY', 'RZ']:
+                theta = gate[1] if input_params is None else input_params[gate[1]]
+                # Scale parameters to prevent gradient explosion
+                theta = np.clip(theta, -np.pi, np.pi)
+                self.register.apply_gate(getattr(QuantumGates, gate[0])(theta), gate[2])
+            else:
+                gate_op = getattr(QuantumGates, gate[0])()
+                if gate[0] == 'CNOT':
+                    self.register.apply_controlled_gate(gate_op, gate[1], gate[2])
+                else:
+                    self.register.apply_gate(gate_op, gate[1])
+        
+        # Return the final state with gradient tracking
+        return Tensor(self.register.state, requires_grad=True)
+
+class QuantumLayer(Module):
+    def __init__(self, num_qubits, num_params):
+        super().__init__()
+        self.circuit = QuantumCircuit(num_qubits)
+        self.params = Tensor(np.random.randn(num_params) * 0.1, requires_grad=True)
+    
+    def build_circuit(self):
+        """Override this method to define the quantum circuit architecture"""
+        raise NotImplementedError
+    
+    def forward(self, x):
+        """Execute quantum circuit with current parameters"""
+        self.build_circuit()
+        return self.circuit.forward(self.params)