iba/__init__.py

from .utils import count_parameters, mark_iba_as_trainable_only
from .XS_llama import IbaXs_LlamaModel, IbaXs_LlamaForCausalLM
from .Xslora import LoraXSLinear, HyperNetXSexp
from .configIBA import MainConfig

# Define what is available when someone uses: from iba import *
__all__ = ["count_parameters", "LoraXSLinear", "HyperNetXSexp", 'IbaXs_LlamaModel', 
           'mark_iba_as_trainable_only', 'IbaXs_LlamaForCausalLM','MainConfig']