update package inits

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

Author: pszemraj

samba_pytorch/__init__.py

@@ -0,0 +1,30 @@
+"""Minimal implementation of Samba by Microsoft in PyTorch."""
+
+from samba_pytorch.config import Config
+from samba_pytorch.samba import GPT, Block, CausalSelfAttention, LLaMAMLP
+from samba_pytorch.tokenizer import Tokenizer
+from samba_pytorch.utils import (
+    chunked_cross_entropy,
+    get_default_supported_precision,
+    lazy_load,
+    num_parameters,
+)
+
+try:
+    from samba_pytorch._version import version as __version__
+except ImportError:
+    __version__ = "0.0.0"
+
+__all__ = [
+    "Config",
+    "GPT",
+    "Block",
+    "CausalSelfAttention",
+    "LLaMAMLP",
+    "Tokenizer",
+    "chunked_cross_entropy",
+    "get_default_supported_precision",
+    "lazy_load",
+    "num_parameters",
+    "__version__",
+]

samba_pytorch/modules/__init__.py

@@ -0,0 +1,24 @@
+"""Core model component modules."""
+
+from samba_pytorch.modules.fused_rotary_embedding import (
+    ApplyRotaryEmb,
+    apply_rotary_emb_func,
+)
+from samba_pytorch.modules.gla import GatedLinearAttention
+from samba_pytorch.modules.mamba_simple import Mamba
+from samba_pytorch.modules.multiscale_retention import MultiScaleRetention
+from samba_pytorch.modules.rmsnorm import FusedRMSNorm, RMSNorm, rms_norm
+from samba_pytorch.modules.rotary import RotaryEmbedding, apply_rotary_emb
+
+__all__ = [
+    "apply_rotary_emb_func",
+    "ApplyRotaryEmb",
+    "GatedLinearAttention",
+    "Mamba",
+    "MultiScaleRetention",
+    "FusedRMSNorm",
+    "RMSNorm",
+    "rms_norm",
+    "apply_rotary_emb",
+    "RotaryEmbedding",
+]