init cleaup

zeta/nn/__init__.py

@@ -1,18 +1,16 @@
 # Attention
 # from zeta.nn.attention import *
-from zeta.nn import attention
+from zeta.nn.attention import *
 
-# architecture
-import zeta.structs as architecture
 
 # embeddings
 # from zeta.nn.embeddings import *
-from zeta.nn import embeddings
+from zeta.nn.embeddings import *
 
 # modules
 # from zeta.nn.modules import *
-from zeta.nn import modules
+from zeta.nn.modules import *
 
 # biases
 # from zeta.nn.biases import *
-from zeta.nn import biases
+from zeta.nn.biases import *

zeta/nn/modules/__init__.py

@@ -57,6 +57,8 @@
 from zeta.nn.modules.itca import IterativeCrossSelfAttention
 from zeta.nn.modules.lang_conv_module import ConvolutionLanguageBlock
 from zeta.nn.modules.h3 import s4d_kernel
+from zeta.nn.modules.h3 import H3Layer
+
 __all__ = [
     "CNNNew",
     "CombinedLinear",
@@ -102,4 +104,5 @@
     "PreNorm",
     "IterativeCrossSelfAttention",
     "ConvolutionLanguageBlock",
+    "H3Layer",
 ]