[FEATS][ ScaledSinusoidalEmbedding ] [ScaleNorm] [ReluSquared]

docs/zeta/nn/attention/local.md

@@ -15,7 +15,8 @@ Key terms:
 ## Class Definition
 
 ```python
-class LocalAttention(nn.Module): ...
+class LocalAttention(nn.Module):
+ ...
 ```
 
 ### Parameters

docs/zeta/nn/attention/mixture_of_attention_ar.md

@@ -32,7 +32,8 @@ class MixtureOfAutoregressiveAttention(nn.Module):
  prenorm: bool = True,
  average_routed: bool = False,
  **kwargs,
- ): ...
+ ):
+ ...
 ```
 
 ### Parameters:
@@ -62,7 +63,8 @@ def forward(
  rotary_emb: Optional[torch.Tensor] = None,
  num_routed_queries: Optional[int] = None,
  num_routed_key_values: Optional[int] = None,
-) -> torch.Tensor: ...
+) -> torch.Tensor:
+ ...
 ```
 
 - `x` (torch.Tensor): Input tensor of shape `(batch_size, sequence_length, dim)`.

docs/zeta/nn/biases/dynamic.md

@@ -15,7 +15,8 @@ Key concepts:
 
 ```python
 class DynamicPositionBias(nn.Module):
- def __init__(self, dim: int, heads: int): ...
+ def __init__(self, dim: int, heads: int):
+ ...
 ```
 
 ### Parameters:

docs/zeta/nn/embeddings/rope.md

@@ -14,7 +14,8 @@ class RotaryEmbedding(nn.Module):
  interpolation_factor=1.0,
  base=10000,
  base_rescale_factor=1.0,
- ): ...
+ ):
+ ...
 ```
 
 ### Parameters
@@ -29,7 +30,8 @@ class RotaryEmbedding(nn.Module):
 ### Method: `forward`
 
 ```python
-def forward(self, seq_len, device): ...
+def forward(self, seq_len, device):
+ ...
 ```
 
 #### Parameters

docs/zeta/nn/modules/token_learner.md

@@ -19,7 +19,8 @@ class TokenLearner(nn.Module):
  ff_mult: int = 2,
  num_output_tokens: int = 8,
  num_layers: int = 2,
- ): ...
+ ):
+ ...
 ```
 
 ### Parameters:
@@ -43,7 +44,8 @@ The forward method of the `TokenLearner` class takes an input tensor `x` and per
 ### Method:
 
 ```python
-def forward(self, x): ...
+def forward(self, x):
+ ...
 ```
 
 ### Parameters:

docs/zeta/nn/modules/visual_expert.md

@@ -32,9 +32,11 @@ class VisualExpert:
  hidden_dim: int,
  dropout: float,
  heads: int,
- ): ...
+ ):
+ ...
 
- def __call__(self, x: torch.Tensor): ...
+ def __call__(self, x: torch.Tensor):
+ ...
 ```
 
 ### Parameters <a name="parameters"></a>

scripts/install_cuda.py

@@ -16,7 +16,7 @@
  "120": "https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda_12.0.1_525.85.12_linux.run",
  "121": "https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run",
  "122": "https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run",
- "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.run",
+ "123": "https://developer.download.nvidia.com/compute/cuda/12.3.2/local_installers/cuda_12.3.2_545.23.08_linux.runbl",
 }
 
 

zeta/nn/embeddings/__init__.py

@@ -27,6 +27,10 @@
  rotate_every_two,
 )
 from zeta.nn.embeddings.yarn import YarnEmbedding
+from zeta.nn.embeddings.scaled_sinusoidal_embeddings import (
+ ScaledSinusoidalEmbedding,
+)
+
 
 __all__ = [
  "AbsolutePositionalEmbedding",
@@ -56,4 +60,5 @@
  "fixed_pos_embedding",
  "duplicate_interleave",
  "VisionEmbedding",
+ "ScaledSinusoidalEmbedding",
 ]

zeta/nn/embeddings/scaled_sinusoidal_embeddings.py

@@ -0,0 +1,47 @@
+import torch
+from torch import nn, Tensor, einsum
+
+from zeta.utils.main import divisible_by
+
+
+class ScaledSinusoidalEmbedding(nn.Module):
+ def __init__(self, dim: int, theta: int = 10000):
+ """
+ Initializes a ScaledSinusoidalEmbedding module.
+
+ Args:
+ dim (int): The dimension of the embedding.
+ theta (int, optional): The scaling factor for the sinusoidal frequencies. Defaults to 10000.
+ """
+ super().__init__()
+ assert divisible_by(dim, 2)
+ self.scale = nn.Parameter(torch.ones(1) * dim**-0.5)
+
+ half_dim = dim // 2
+ freq_seq = torch.arange(half_dim).float() / half_dim
+ inv_freq = theta**-freq_seq
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, x: Tensor, pos=None, seq_start_pos=None):
+ """
+ Forward pass of the ScaledSinusoidalEmbedding module.
+
+ Args:
+ x (Tensor): The input tensor.
+ pos (Tensor, optional): The position tensor. Defaults to None.
+ seq_start_pos (Tensor, optional): The starting position tensor for sequences. Defaults to None.
+
+ Returns:
+ Tensor: The embedded tensor.
+ """
+ sq, device = x.shape[1], x.device
+
+ if pos is not None:
+ pos = torch.arange(sq, device=device)
+
+ if seq_start_pos is not None:
+ pos = pos - seq_start_pos[..., None]
+
+ emb = einsum("i, j -> i j", pos, self.inv_freq)
+ emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+ return emb * self.scale

zeta/nn/modules/__init__.py

@@ -188,6 +188,8 @@
 from zeta.nn.modules.ws_conv2d import WSConv2d
 from zeta.nn.modules.yolo import yolo
 from zeta.nn.modules.palo_ldp import PaloLDP
+from zeta.nn.modules.relu_squared import ReluSquared
+from zeta.nn.modules.scale_norm import ScaleNorm
 
 # from zeta.nn.modules.g_shard_moe import (
 # Top1Gate,
@@ -386,4 +388,6 @@
  "DynamicInputChannels",
  "OutputDecoders",
  "PaloLDP",
+ "ReluSquared",
+ "ScaleNorm",
 ]

zeta/nn/modules/relu_squared.py

@@ -0,0 +1,17 @@
+from torch import nn
+import torch.nn.functional as F
+
+
+class ReluSquared(nn.Module):
+ """
+ Applies the ReLU activation function and squares the output.
+
+ Args:
+ x (torch.Tensor): Input tensor.
+
+ Returns:
+ torch.Tensor: Output tensor after applying ReLU and squaring the result.
+ """
+
+ def forward(self, x):
+ return F.relu(x) ** 2

zeta/nn/modules/scale_norm.py

@@ -0,0 +1,35 @@
+import torch
+from torch import nn, Tensor
+
+
+class ScaleNorm(nn.Module):
+ """
+ Applies scale normalization to the input tensor along the last dimension.
+
+ Args:
+ dim (int): The dimension of the input tensor.
+ eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-5.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ eps: float = 1e-5,
+ ):
+ super().__init__()
+ self.eps = eps
+
+ self.g = nn.Parameter(torch.ones(1) * (dim**-0.5))
+
+ def forward(self, x: Tensor):
+ """
+ Applies scale normalization to the input tensor.
+
+ Args:
+ x (Tensor): The input tensor.
+
+ Returns:
+ Tensor: The scale-normalized tensor.
+ """
+ norm = torch.norm(x, dim=-1, keepdim=True)
+ return x / norm.clamp(min=self.eps) + self.g