# mamba_ssm package init — minimal override to avoid broken selective_scan_cuda.so
# ABI mismatch with the base image's libtorch.
#
# The upstream __init__.py eagerly imports selective_scan_cuda which fails on
# pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel (undefined c10::Warning ctor
# symbol). We only need Mamba3 (grafted from main, pure-Triton), so we skip
# all compiled-CUDA imports here and let Mamba3 load directly.

__version__ = "2.3.1+feather-graft"

# selective_scan_fn / mamba_inner_fn are shimmed to None — they are NOT used
# by the Feather training path (which is Mamba3-only). If any import path
# hits this, it will get a clear AttributeError instead of an obscure ImportError.
selective_scan_fn = None
mamba_inner_fn = None

# --- triton API compatibility shims -----------------------------------------
# Version matrix is hostile: torch 2.6 pins triton==3.2.0 because torch._inductor
# imports AttrsDescriptor from triton.compiler.compiler — removed in triton 3.4+.
# Grafted Mamba3 (from mamba-ssm main) needs triton.set_allocator and
# tl.make_tensor_descriptor, both added in triton 3.3+. No single triton version
# satisfies both simultaneously. We run on triton 3.5.1 (latest, has both mamba3
# APIs) and shim AttrsDescriptor as a stub dataclass for torch._inductor. The
# stub is never actually invoked at runtime because the codebase does not use
# torch.compile — but importing torch._inductor.* still requires the symbol to
# exist at module load time.
import triton as _triton  # noqa: E402
if not hasattr(_triton, "set_allocator"):
    def _noop_set_allocator(_fn):  # pragma: no cover
        return None
    _triton.set_allocator = _noop_set_allocator

import triton.compiler.compiler as _tcc  # noqa: E402
if not hasattr(_tcc, "AttrsDescriptor"):
    class _AttrsDescriptorShim:
        """Stub for torch._inductor compatibility on triton >= 3.4.
        torch._inductor.runtime.hints imports this at module load but the
        constructor is only called inside torch.compile paths. Accept any
        args/kwargs so the import itself succeeds."""
        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs

        @classmethod
        def from_hints(cls, *args, **kwargs):
            return cls(*args, **kwargs)

    _tcc.AttrsDescriptor = _AttrsDescriptorShim

# triton_key: removed in triton 3.5, used by torch._inductor.codecache for
# FxGraphCache key derivation. Return a stable string so caching still works.
if not hasattr(_tcc, "triton_key"):
    def _triton_key_shim():
        import triton as _t
        return f"triton-{_t.__version__}-shim"
    _tcc.triton_key = _triton_key_shim

# Suppress torch.compile/_dynamo errors globally — we don't rely on torch.compile
# for performance in this codebase (Muon + mamba3 CUDA kernels already fused),
# so fall back to eager on any dynamo failure rather than crashing. This is
# defense-in-depth against further triton API drift.
try:
    import torch._dynamo  # noqa: F401 — triggers dynamo module init
    torch._dynamo.config.suppress_errors = True
except Exception:  # pragma: no cover
    pass

# Expose Mamba3 at top level to match `from mamba_ssm import Mamba3`.
from mamba_ssm.modules.mamba3 import Mamba3  # noqa: E402