| """ | |
| BURT-IMMA: Bi-encoder Unified Retrieval Transformer with | |
| Interferometric Matrix Memory Architecture | |
| License: BSL-1.1 | |
| Contact: jessica@collectivekitty.com | |
| """ | |
| __version__ = "0.1.0" | |
| # Try to import CUDA-accelerated kernels, fall back to pure PyTorch | |
| _CUDA_AVAILABLE = False | |
| try: | |
| import _burt_imma_cuda | |
| _CUDA_AVAILABLE = True | |
| except ImportError: | |
| _burt_imma_cuda = None | |
| # Import kernel wrappers (these handle CUDA/CPU fallback internally) | |
| from .kernels import ( | |
| constrained_softmax, | |
| cifg_update, | |
| batched_cifg_update, | |
| sparse_moe_dispatch, | |
| biencoder_attention, | |
| attention_softmax, | |
| ) | |
| from .utils import ( | |
| spectral_norm, | |
| entropy, | |
| check_huntington, | |
| load_config, | |
| set_seed, | |
| Timer, | |
| ) | |
| __all__ = [ | |
| # Kernel functions | |
| "constrained_softmax", | |
| "cifg_update", | |
| "batched_cifg_update", | |
| "sparse_moe_dispatch", | |
| "biencoder_attention", | |
| "attention_softmax", | |
| # Utilities | |
| "spectral_norm", | |
| "entropy", | |
| "check_huntington", | |
| "load_config", | |
| "set_seed", | |
| "Timer", | |
| # Meta | |
| "__version__", | |
| "_CUDA_AVAILABLE", | |
| ] | |