""" 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
from zeta.nn.modules.combined_linear import CombinedLinear
from zeta.nn.modules.convnet import ConvNet
from zeta.nn.modules.droppath import DropPath
from zeta.nn.modules.dynamic_module import DynamicModule
from zeta.nn.modules.ether import Ether
from zeta.nn.modules.exo import Exo
from zeta.nn.modules.fast_text import FastTextNew
from zeta.nn.modules.feedforward import FeedForward
from zeta.nn.modules.feedforward_network import FeedForwardNetwork
from zeta.nn.modules.flexible_mlp import CustomMLP
from zeta.nn.modules.h3 import H3Layer
from zeta.nn.modules.itca import IterativeCrossSelfAttention
from zeta.nn.modules.lang_conv_module import ConvolutionLanguageBlock
from zeta.nn.modules.layernorm import LayerNorm, l2norm
from zeta.nn.modules.leaky_relu import LeakyRELU
from zeta.nn.modules.log_ff import LogFF
from zeta.nn.modules.lora import Lora
from zeta.nn.modules.mbconv import (
    DropSample,
    SqueezeExcitation,
    MBConvResidual,
    MBConv,
)
from zeta.nn.modules.mlp import MLP
from zeta.nn.modules.mlp_mixer import MLPBlock, MixerBlock, MLPMixer
from zeta.nn.modules.nebula import Nebula
from zeta.nn.modules.polymorphic_activation import PolymorphicActivation
from zeta.nn.modules.polymorphic_neuron import PolymorphicNeuronLayer
from zeta.nn.modules.prenorm import PreNorm
from zeta.nn.modules.pulsar import Pulsar
from zeta.nn.modules.residual import Residual
from zeta.nn.modules.resnet import ResNet
from zeta.nn.modules.rms_norm import RMSNorm
from zeta.nn.modules.rnn_nlp import RNNL
from zeta.nn.modules.shufflenet import ShuffleNet
from zeta.nn.modules.sig_lip import SigLipLoss
from zeta.nn.modules.simple_attention import simple_attention
from zeta.nn.modules.simple_feedforward import SimpleFeedForward
from zeta.nn.modules.simple_res_block import SimpleResBlock
from zeta.nn.modules.skipconnection import SkipConnection
from zeta.nn.modules.spatial_transformer import SpatialTransformer
from zeta.nn.modules.subln import SubLN
from zeta.nn.modules.super_resolution import SuperResolutionNet
from zeta.nn.modules.time_up_sample import TimeUpSample2x
from zeta.nn.modules.token_learner import TokenLearner
from zeta.nn.modules.unet import Unet
from zeta.nn.modules.video_autoencoder import CausalConv3d
from zeta.nn.modules.visual_expert import VisualExpert
from zeta.nn.modules.yolo import yolo
from zeta.nn.modules.swiglu import SwiGLU, SwiGLUStacked
from zeta.nn.modules.img_patch_embed import ImgPatchEmbed
from zeta.nn.modules.dense_connect import DenseBlock
from zeta.nn.modules.highway_layer import HighwayLayer
from zeta.nn.modules.multi_scale_block import MultiScaleBlock
from zeta.nn.modules.feedback_block import FeedbackBlock
from zeta.nn.modules.dual_path_block import DualPathBlock
from zeta.nn.modules.recursive_block import RecursiveBlock
from zeta.nn.modules._activations import (
    PytorchGELUTanh,
    NewGELUActivation,
    GELUActivation,
    FastGELUActivation,
    QuickGELUActivation,
    ClippedGELUActivation,
    AccurateGELUActivation,
    MishActivation,
    LinearActivation,
    LaplaceActivation,
    ReLUSquaredActivation,
)


from zeta.nn.modules.triple_skip import TripleSkipBlock
from zeta.nn.modules.dynamic_routing_block import DynamicRoutingBlock
from zeta.nn.modules.gated_residual_block import GatedResidualBlock
from zeta.nn.modules.stochastic_depth import StochasticSkipBlocK

#######
from zeta.nn.modules.quantized_layernorm import QuantizedLN
from zeta.nn.modules.slerp_model_merger import SLERPModelMerger
from zeta.nn.modules.avg_model_merger import AverageModelMerger
from zeta.nn.modules.adaptive_rmsnorm import AdaptiveRMSNorm

######
from zeta.nn.modules.simple_mamba import MambaBlock, Mamba
from zeta.nn.modules.laser import Laser
from zeta.nn.modules.fused_gelu_dense import FusedDenseGELUDense
from zeta.nn.modules.fused_dropout_layernom import FusedDropoutLayerNorm
from zeta.nn.modules.conv_mlp import Conv2DFeedforward
from zeta.nn.modules.ws_conv2d import WSConv2d
from zeta.nn.modules.stoch_depth import StochDepth
from zeta.nn.modules.nfn_stem import NFNStem
from zeta.nn.modules.film import Film
from zeta.nn.modules.video_to_tensor import video_to_tensor, video_to_tensor_vr
from zeta.nn.modules.proj_then_softmax import FusedProjSoftmax
from zeta.nn.modules.top_n_gating import TopNGating
from zeta.nn.modules.moe_router import MoERouter
from zeta.nn.modules.perceiver_layer import PerceiverLayer
from zeta.nn.modules.u_mamba import UMambaBlock
from zeta.nn.modules.audio_to_text import audio_to_text
from zeta.nn.modules.patch_video import patch_video
from zeta.nn.modules.image_to_text import img_to_text
from zeta.nn.modules.video_to_text import video_to_text
from zeta.nn.modules.pyro import hyper_optimize
from zeta.nn.modules.vit_denoiser import (
    to_patch_embedding,
    posemb_sincos_2d,
    VisionAttention,
    VitTransformerBlock,
)
from zeta.nn.modules.v_layernorm import VLayerNorm
from zeta.nn.modules.parallel_wrapper import Parallel
from zeta.nn.modules.v_pool import DepthWiseConv2d, Pool
from zeta.nn.modules.moe import MixtureOfExperts
from zeta.nn.modules.flex_conv import FlexiConv
from zeta.nn.modules.mm_layernorm import MMLayerNorm
from zeta.nn.modules.fusion_ffn import MMFusionFFN
from zeta.nn.modules.norm_utils import PostNorm
from zeta.nn.modules.mm_mamba_block import MultiModalMambaBlock
from zeta.nn.modules.p_scan import PScan, pscan
from zeta.nn.modules.ssm import selective_scan, selective_scan_seq, SSM
from zeta.nn.modules.film_conditioning import FilmConditioning
from zeta.nn.modules.qkv_norm import qkv_norm, qk_norm


####
from zeta.nn.modules.space_time_unet import (
    FeedForwardV,
    ContinuousPositionBias,
    PseudoConv3d,
    SpatioTemporalAttention,
    ResnetBlock,
    Downsample,
    Upsample,
    SpaceTimeUnet,
)
from zeta.nn.modules.patch_img import patch_img
from zeta.nn.modules.mm_ops import threed_to_text, text_to_twod
from zeta.nn.modules.fused_dropout_add import (
    jit_dropout_add,
    fused_dropout_add,
    jit_bias_dropout_add,
    fused_bias_dropout_add,
)
from zeta.nn.modules.blockdiag_butterfly import (
    blockdiag_butterfly_multiply_reference,
    BlockdiagButterflyMultiply,
    blockdiag_weight_to_dense_weight,
    blockdiag_multiply_reference,
    BlockdiagMultiply,
    fftconv_ref,
    mul_sum,
    Sin,
    StructuredLinear,
)

from zeta.nn.modules.block_butterfly_mlp import (
    BlockButterflyLinear,
    BlockMLP,
)

from zeta.nn.modules.gill_mapper import GILLMapper
from zeta.nn.modules.add_norm import add_norm
from zeta.nn.modules.to_logits import to_logits
from zeta.nn.modules.cross_modal_reparametization import (
    CrossModalReparamLinear,
    cross_modal_ffn,
    build_cross_modal_reparam_linear,
    change_original_linear_to_reparam,
    reparameterize_aux_into_target_model,
    CrossModalReParametrization,
)
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
# from zeta.nn.modules.scaled_sinusoidal import ScaledSinuosidalEmbedding
# from zeta.nn.modules.scale import Scale
# from zeta.nn.modules.scalenorm import ScaleNorm
# from zeta.nn.modules.simple_rmsnorm import SimpleRMSNorm
# from zeta.nn.modules.gru_gating import GRUGating
# from zeta.nn.modules.shift_tokens import ShiftTokens
# from zeta.nn.modules.swarmalator import simulate_swarmalators
# from zeta.nn.modules.transformations import image_transform
# from zeta.nn.modules.squeeze_excitation import SqueezeExcitation
# from zeta.nn.modules.clex import Clex

__all__ = [
    "CNNNew",
    "CombinedLinear",
    "ConvNet",
    "DropPath",
    "DynamicModule",
    "Exo",
    "FastTextNew",
    "FeedForwardNetwork",
    "LayerNorm",
    "l2norm",
    "Lora",
    "MBConv",
    "MLP",
    "Pulsar",
    "Residual",
    "ResNet",
    "RMSNorm",
    "RNNL",
    "ShuffleNet",
    "simple_attention",
    "SpatialTransformer",
    "SubLN",
    "SuperResolutionNet",
    "TokenLearner",
    "yolo",
    "Ether",
    "Nebula",
    "AdaptiveConv3DMod",
    "TimeUpSample2x",
    "CausalConv3d",
    "SimpleResBlock",
    "SigLipLoss",
    "SimpleFeedForward",
    "Unet",
    "VisualExpert",
    "FeedForward",
    "SkipConnection",
    "LogFF",
    "PolymorphicNeuronLayer",
    "CustomMLP",
    "PolymorphicActivation",
    "PreNorm",
    "IterativeCrossSelfAttention",
    "ConvolutionLanguageBlock",
    "H3Layer",
    "MLPMixer",
    "LeakyRELU",
    "AdaptiveLayerNorm",
    "SwiGLU",
    "SwiGLUStacked",
    "ImgPatchEmbed",
    "DenseBlock",
    "HighwayLayer",
    "MultiScaleBlock",
    "FeedbackBlock",
    "DualPathBlock",
    "RecursiveBlock",
    "PytorchGELUTanh",
    "NewGELUActivation",
    "GELUActivation",
    "FastGELUActivation",
    "QuickGELUActivation",
    "ClippedGELUActivation",
    "AccurateGELUActivation",
    "MishActivation",
    "LinearActivation",
    "LaplaceActivation",
    "ReLUSquaredActivation",
    "TripleSkipBlock",
    "DynamicRoutingBlock",
    "GatedResidualBlock",
    "StochasticSkipBlocK",
    "QuantizedLN",
    "SLERPModelMerger",
    "AverageModelMerger",
    "AdaptiveRMSNorm",
    "MambaBlock",
    "Mamba",
    "Laser",
    "FusedDenseGELUDense",
    "FusedDropoutLayerNorm",
    "Conv2DFeedforward",
    "MLPBlock",
    "MixerBlock",
    "WSConv2d",
    "StochDepth",
    "NFNStem",
    "Film",
    "DropSample",
    "SqueezeExcitation",
    "MBConvResidual",
    "video_to_tensor",
    "video_to_tensor_vr",
    "FusedProjSoftmax",
    "TopNGating",
    "MoERouter",
    "PerceiverLayer",
    "UMambaBlock",
    "audio_to_text",
    "patch_video",
    "img_to_text",
    "video_to_text",
    "hyper_optimize",
    "to_patch_embedding",
    "posemb_sincos_2d",
    "VisionAttention",
    "VitTransformerBlock",
    "VLayerNorm",
    "Parallel",
    "DepthWiseConv2d",
    "Pool",
    "MixtureOfExperts",
    "FlexiConv",
    "MMLayerNorm",
    "MMFusionFFN",
    "PostNorm",
    "MultiModalMambaBlock",
    "PScan",
    "pscan",
    "selective_scan",
    "selective_scan_seq",
    "SSM",
    "FilmConditioning",
    "qkv_norm",
    "qk_norm",
    "FeedForwardV",
    "ContinuousPositionBias",
    "PseudoConv3d",
    "SpatioTemporalAttention",
    "ResnetBlock",
    "Downsample",
    "Upsample",
    "SpaceTimeUnet",
    "patch_img",
    "threed_to_text",
    "text_to_twod",
    "jit_dropout_add",
    "fused_dropout_add",
    "jit_bias_dropout_add",
    "fused_bias_dropout_add",
    "blockdiag_butterfly_multiply_reference",
    "BlockdiagButterflyMultiply",
    "blockdiag_weight_to_dense_weight",
    "blockdiag_multiply_reference",
    "BlockdiagMultiply",
    "fftconv_ref",
    "mul_sum",
    "Sin",
    "StructuredLinear",
    "BlockButterflyLinear",
    "BlockMLP",
    "GILLMapper",
    "add_norm",
    "to_logits",
    "CrossModalReParametrization",
    "CrossModalReparamLinear",
    "cross_modal_ffn",
    "build_cross_modal_reparam_linear",
    "change_original_linear_to_reparam",
    "reparameterize_aux_into_target_model",
    "QFormer",
    "MLPProjectionFusion",
    "NormalizationFractral",
]
