[CLEANUP]

docs/zeta/nn/modules/dynamicroutingblock.md

@@ -74,7 +74,7 @@ drb = DynamicRoutingBlock(sb1, sb2, routing_module)
 The input can be passed to this block to yield the output:
 
 ```python
-x = torch.randn(10, 5)
+x = torch.randn(3, 5)
 y = drb(x)
 ```
 In the process, the dynamic routing block has learned to route between `sb1` and `sb2` depending on `routing_module`'s weights, allowing the module to discover which sub-block is more 'helpful' for any given input.

tests/cloud/test_main.py

@@ -1,3 +1,5 @@
+"""Test cases for the main module of the cloud package."""
+
 import pytest
 from unittest.mock import MagicMock, patch
 from zeta.cloud.main import zetacloud

tests/nn/attentions/test_attend.py

@@ -1,3 +1,5 @@
+""" Test cases for the Attend module. """
+
 import torch
 from zeta.nn.attention.attend import Attend
 
@@ -121,6 +123,21 @@ def test_attend_flash_attention():
  assert out.shape == (1, 8, 32, 64)
 
 
+# Test case for configuring flash attention
+def test_flash_attention():
+ import torch
+ from zeta.nn import FlashAttention
+
+ q = torch.randn(2, 4, 6, 8)
+ k = torch.randn(2, 4, 10, 8)
+ v = torch.randn(2, 4, 10, 8)
+
+ attention = FlashAttention(causal=False, dropout=0.1, flash=True)
+ output = attention(q, k, v)
+
+ assert output.shape == (2, 4, 6, 8)
+
+
 # Test case for gradient checking using torch.autograd.gradcheck
 def test_attend_gradient_check():
  attend = Attend()

zeta/cloud/__init__.py

@@ -1,3 +1,4 @@
+""" init file for cloud module """
 from zeta.cloud.sky_api import SkyInterface
 from zeta.cloud.main import zetacloud
 

zeta/cloud/main.py

@@ -1,3 +1,5 @@
+"""Cloud """
+
 import logging
 from typing import Any
 

zeta/cloud/sky_api.py

@@ -1,3 +1,7 @@
+""" sky_api module """
+""" This module provides a simplified interface for launching, executing, 
+stopping, starting, and tearing down clusters. """
+
 from typing import List
 
 import sky

zeta/nn/__init__.py

@@ -1,3 +1,4 @@
+""" Neural network modules. zeta/nn """
 from zeta.nn.attention import * # noqa: F403
 from zeta.nn.embeddings import * # noqa: F403
 from zeta.nn.modules import * # noqa: F403

zeta/nn/attention/__init__.py

@@ -1,4 +1,4 @@
-"""Zeta Halo"""
+"""Zeta Attention init file"""
 from zeta.nn.attention.attend import Attend, Intermediates
 from zeta.nn.attention.cross_attn_images import MultiModalCrossAttention
 from zeta.nn.attention.flash_attention import FlashAttention

zeta/nn/modules/__init__.py

@@ -1,3 +1,4 @@
+""" init file for nn modules """
 from zeta.nn.modules.adaptive_conv import AdaptiveConv3DMod
 from zeta.nn.modules.adaptive_layernorm import AdaptiveLayerNorm
 from zeta.nn.modules.cnn_text import CNNNew
@@ -175,7 +176,7 @@
 )
 from zeta.nn.modules.qformer import QFormer
 from zeta.nn.modules.poly_expert_fusion_network import MLPProjectionFusion
-
+from zeta.nn.modules.norm_fractorals import NormalizationFractral
 
 # from zeta.nn.modules.img_reshape import image_reshape
 # from zeta.nn.modules.flatten_features import flatten_features
@@ -353,4 +354,5 @@
  "reparameterize_aux_into_target_model",
  "QFormer",
  "MLPProjectionFusion",
+ "NormalizationFractral",
 ]

zeta/nn/modules/norm_fractorals.py

@@ -0,0 +1,52 @@
+from torch import nn
+
+
+class NormalizationFractral(nn.Module):
+ """
+ A module that performs normalization using fractal layers.
+
+ Args:
+ dim (int): The input dimension.
+ eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-8.
+ fi (int, optional): The number of fractal layers. Default is 4.
+ *args: Variable length argument list.
+ **kwargs: Arbitrary keyword arguments.
+
+ Attributes:
+ eps (float): A small value added to the denominator for numerical stability.
+ fi (int): The number of fractal layers.
+ norm (nn.LayerNorm): The initial normalization layer.
+ norm_i (nn.LayerNorm): Fractal normalization layers.
+
+ """
+
+ def __init__(
+ self, dim: int, eps=1e-8, fi: int = 4, *args, **kwargs # Fractal index
+ ):
+ super(NormalizationFractral, self).__init__(*args, **kwargs)
+ self.eps = eps
+ self.fi = fi
+
+ self.norm = nn.LayerNorm(dim)
+
+ for i in range(fi):
+ setattr(self, f"norm_{i}", nn.LayerNorm(dim))
+
+ def forward(self, x):
+ """
+ Forward pass of the NormalizationFractral module.
+
+ Args:
+ x (torch.Tensor): The input tensor.
+
+ Returns:
+ torch.Tensor: The normalized output tensor.
+
+ """
+ x = self.norm(x)
+
+ for i in range(self.fi):
+ norm = getattr(self, f"norm_{i}")
+ x = norm(x)
+
+ return x

zeta/nn/modules/qformer.py

@@ -1,10 +1,18 @@
+""" QFormer module for processing text and image inputs. """
+
 from einops import rearrange, reduce
 from torch import Tensor, nn
 
+<<<<<<< HEAD
 from zeta.nn.attention.multiquery_attention import (
  MultiQueryAttention,
 )
 from zeta.nn.modules.simple_feedforward import SimpleFeedForward
+=======
+from zeta.nn.attention.multiquery_attention import MultiQueryAttention
+from zeta.nn.modules import SimpleFeedForward
+
+>>>>>>> 2e50e6fbb49a66ed3ef6cf19426fbbd191ca61aa
 from zeta.nn.attention.cross_attention import CrossAttention