                for pth in dir.glob(lib_pattern):
                    if pth.is_file() and not pth.is_symlink():
                        yield pth
        except (OSError, PermissionError):
            pass

def is_relevant_candidate_env_var(env_var: str, value: str) -> bool:
    return (
        env_var in CUDART_PATH_PREFERRED_ENVVARS  # is a preferred location
        or (
            os.sep in value  # might contain a path
            and env_var not in CUDART_PATH_IGNORED_ENVVARS  # not ignored
            and "CONDA" not in env_var  # not another conda envvar
            and "BASH_FUNC" not in env_var  # not a bash function defined via envvar
            and "\n" not in value  # likely e.g. a script or something?
        )
    )

def get_potentially_lib_path_containing_env_vars() -> dict[str, str]:
    return {env_var: value for env_var, value in os.environ.items() if is_relevant_candidate_env_var(env_var, value)}

def find_cudart_libraries() -> Iterator[Path]:
    """
    Searches for a cuda installations, in the following order of priority:
        1. active conda env
        2. LD_LIBRARY_PATH
        3. any other env vars, while ignoring those that
            - are known to be unrelated
            - don't contain the path separator `/`

    If multiple libraries are found in part 3, we optimistically try one,
    while giving a warning message.
    """
    candidate_env_vars = get_potentially_lib_path_containing_env_vars()

    for envvar in CUDART_PATH_PREFERRED_ENVVARS:
        if envvar in candidate_env_vars:
            directory = candidate_env_vars[envvar]
            yield from find_cuda_libraries_in_path_list(directory)


            pass

def is_relevant_candidate_env_var(env_var: str, value: str) -> bool:
    return (
        env_var in CUDART_PATH_PREFERRED_ENVVARS  # is a preferred location
        or (
            os.sep in value  # might contain a path
            and env_var not in CUDART_PATH_IGNORED_ENVVARS  # not ignored
            and "CONDA" not in env_var  # not another conda envvar
            and "BASH_FUNC" not in env_var  # not a bash function defined via envvar
            and "\n" not in value  # likely e.g. a script or something?
        )
    )

def get_potentially_lib_path_containing_env_vars() -> dict[str, str]:
    return {env_var: value for env_var, value in os.environ.items() if is_relevant_candidate_env_var(env_var, value)}

def find_cudart_libraries() -> Iterator[Path]:
    """
    Searches for a cuda installations, in the following order of priority:
        1. active conda env
        2. LD_LIBRARY_PATH
        3. any other env vars, while ignoring those that
            - are known to be unrelated
            - don't contain the path separator `/`

    If multiple libraries are found in part 3, we optimistically try one,
    while giving a warning message.
    """
    candidate_env_vars = get_potentially_lib_path_containing_env_vars()

    for envvar in CUDART_PATH_PREFERRED_ENVVARS:
        if envvar in candidate_env_vars:
            directory = candidate_env_vars[envvar]
            yield from find_cuda_libraries_in_path_list(directory)
            candidate_env_vars.pop(envvar)

    for env_var, value in candidate_env_vars.items():
        yield from find_cuda_libraries_in_path_list(value)

def _print_cuda_diagnostics(cuda_specs: CUDASpecs) -> None:
    print(
        f"PyTorch settings found: CUDA_VERSION={cuda_specs.cuda_version_string}, "
        f"Highest Compute Capability: {cuda_specs.highest_compute_capability}.",
    )

    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
            Library not found: {binary_path}. Maybe you need to compile it from source?
            """,
        )

    # 7.5 is the minimum CC for int8 tensor cores
    if not cuda_specs.has_imma:
        print_dedented(
            """


        2. LD_LIBRARY_PATH
        3. any other env vars, while ignoring those that
            - are known to be unrelated
            - don't contain the path separator `/`

    If multiple libraries are found in part 3, we optimistically try one,
    while giving a warning message.
    """
    candidate_env_vars = get_potentially_lib_path_containing_env_vars()

    for envvar in CUDART_PATH_PREFERRED_ENVVARS:
        if envvar in candidate_env_vars:
            directory = candidate_env_vars[envvar]
            yield from find_cuda_libraries_in_path_list(directory)
            candidate_env_vars.pop(envvar)

    for env_var, value in candidate_env_vars.items():
        yield from find_cuda_libraries_in_path_list(value)

def _print_cuda_diagnostics(cuda_specs: CUDASpecs) -> None:
    print(
        f"PyTorch settings found: CUDA_VERSION={cuda_specs.cuda_version_string}, "
        f"Highest Compute Capability: {cuda_specs.highest_compute_capability}.",
    )

    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
            Library not found: {binary_path}. Maybe you need to compile it from source?
            """,
        )

    # 7.5 is the minimum CC for int8 tensor cores
    if not cuda_specs.has_imma:
        print_dedented(
            """
            WARNING: Compute capability < 7.5 detected! Only slow 8-bit matmul is supported for your GPU!
            If you run into issues with 8-bit matmul, you can try 4-bit quantization:
            https://huggingface.co/blog/4bit-transformers-bitsandbytes
            """,
        )

def _print_hip_diagnostics(cuda_specs: CUDASpecs) -> None:
    print(f"PyTorch settings found: ROCM_VERSION={cuda_specs.cuda_version_string}")

    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
        Library not found: {binary_path}.
        Maybe you need to compile it from source? If you compiled from source, check that ROCm version
        in PyTorch Settings matches your ROCm install. If not, reinstall PyTorch for your ROCm version
        and rebuild bitsandbytes.
        """,
        )

    hip_major, hip_minor = cuda_specs.cuda_version_tuple
    if (hip_major, hip_minor) < (6, 1):
        print_dedented(
            """


    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
            Library not found: {binary_path}. Maybe you need to compile it from source?
            """,
        )

    # 7.5 is the minimum CC for int8 tensor cores
    if not cuda_specs.has_imma:
        print_dedented(
            """
            WARNING: Compute capability < 7.5 detected! Only slow 8-bit matmul is supported for your GPU!
            If you run into issues with 8-bit matmul, you can try 4-bit quantization:
            https://huggingface.co/blog/4bit-transformers-bitsandbytes
            """,
        )

def _print_hip_diagnostics(cuda_specs: CUDASpecs) -> None:
    print(f"PyTorch settings found: ROCM_VERSION={cuda_specs.cuda_version_string}")

    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
        Library not found: {binary_path}.
        Maybe you need to compile it from source? If you compiled from source, check that ROCm version
        in PyTorch Settings matches your ROCm install. If not, reinstall PyTorch for your ROCm version
        and rebuild bitsandbytes.
        """,
        )

    hip_major, hip_minor = cuda_specs.cuda_version_tuple
    if (hip_major, hip_minor) < (6, 1):
        print_dedented(
            """
            WARNING: bitsandbytes is fully supported only from ROCm 6.1.
            """,
        )

def print_diagnostics(cuda_specs: CUDASpecs) -> None:
    if HIP_ENVIRONMENT:
        _print_hip_diagnostics(cuda_specs)
    else:
        _print_cuda_diagnostics(cuda_specs)

def _print_cuda_runtime_diagnostics() -> None:
    cudart_paths = list(find_cudart_libraries())
    if not cudart_paths:
        print("CUDA SETUP: WARNING! CUDA runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate CUDA runtime files (see below).

            We select the PyTorch default CUDA runtime, which is {torch.version.cuda},
            but this might mismatch with the CUDA version that is needed for bitsandbytes.


    binary_path = get_cuda_bnb_library_path(cuda_specs)
    if not binary_path.exists():
        print_dedented(
            f"""
        Library not found: {binary_path}.
        Maybe you need to compile it from source? If you compiled from source, check that ROCm version
        in PyTorch Settings matches your ROCm install. If not, reinstall PyTorch for your ROCm version
        and rebuild bitsandbytes.
        """,
        )

    hip_major, hip_minor = cuda_specs.cuda_version_tuple
    if (hip_major, hip_minor) < (6, 1):
        print_dedented(
            """
            WARNING: bitsandbytes is fully supported only from ROCm 6.1.
            """,
        )

def print_diagnostics(cuda_specs: CUDASpecs) -> None:
    if HIP_ENVIRONMENT:
        _print_hip_diagnostics(cuda_specs)
    else:
        _print_cuda_diagnostics(cuda_specs)

def _print_cuda_runtime_diagnostics() -> None:
    cudart_paths = list(find_cudart_libraries())
    if not cudart_paths:
        print("CUDA SETUP: WARNING! CUDA runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate CUDA runtime files (see below).

            We select the PyTorch default CUDA runtime, which is {torch.version.cuda},
            but this might mismatch with the CUDA version that is needed for bitsandbytes.
            To override this behavior set the `BNB_CUDA_VERSION=<version string, e.g. 122>` environmental variable.

            For example, if you want to use the CUDA version 122,
                BNB_CUDA_VERSION=122 python ...

            OR set the environmental variable in your .bashrc:
                export BNB_CUDA_VERSION=122


        and rebuild bitsandbytes.
        """,
        )

    hip_major, hip_minor = cuda_specs.cuda_version_tuple
    if (hip_major, hip_minor) < (6, 1):
        print_dedented(
            """
            WARNING: bitsandbytes is fully supported only from ROCm 6.1.
            """,
        )

def print_diagnostics(cuda_specs: CUDASpecs) -> None:
    if HIP_ENVIRONMENT:
        _print_hip_diagnostics(cuda_specs)
    else:
        _print_cuda_diagnostics(cuda_specs)

def _print_cuda_runtime_diagnostics() -> None:
    cudart_paths = list(find_cudart_libraries())
    if not cudart_paths:
        print("CUDA SETUP: WARNING! CUDA runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate CUDA runtime files (see below).

            We select the PyTorch default CUDA runtime, which is {torch.version.cuda},
            but this might mismatch with the CUDA version that is needed for bitsandbytes.
            To override this behavior set the `BNB_CUDA_VERSION=<version string, e.g. 122>` environmental variable.

            For example, if you want to use the CUDA version 122,
                BNB_CUDA_VERSION=122 python ...

            OR set the environmental variable in your .bashrc:
                export BNB_CUDA_VERSION=122

            In the case of a manual override, make sure you set LD_LIBRARY_PATH, e.g.
            export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.2,
            """,
        )
        for pth in cudart_paths:
            print(f"* Found CUDA runtime at: {pth}")

def _print_hip_runtime_diagnostics() -> None:
    cudart_paths = list(find_cudart_libraries())
    if not cudart_paths:
        print("WARNING! ROCm runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate ROCm runtime files (see below).

            We select the PyTorch default ROCm runtime, which is {torch.version.hip},
            but this might mismatch with the ROCm version that is needed for bitsandbytes.

            To resolve it, install PyTorch built for the ROCm version you want to use

            and set LD_LIBRARY_PATH to your ROCm install path, e.g.
            export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm-6.1.2/lib,
            """,
        )


            Found duplicate CUDA runtime files (see below).

            We select the PyTorch default CUDA runtime, which is {torch.version.cuda},
            but this might mismatch with the CUDA version that is needed for bitsandbytes.
            To override this behavior set the `BNB_CUDA_VERSION=<version string, e.g. 122>` environmental variable.

            For example, if you want to use the CUDA version 122,
                BNB_CUDA_VERSION=122 python ...

            OR set the environmental variable in your .bashrc:
                export BNB_CUDA_VERSION=122

            In the case of a manual override, make sure you set LD_LIBRARY_PATH, e.g.
            export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.2,
            """,
        )
        for pth in cudart_paths:
            print(f"* Found CUDA runtime at: {pth}")

def _print_hip_runtime_diagnostics() -> None:
    cudart_paths = list(find_cudart_libraries())
    if not cudart_paths:
        print("WARNING! ROCm runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate ROCm runtime files (see below).

            We select the PyTorch default ROCm runtime, which is {torch.version.hip},
            but this might mismatch with the ROCm version that is needed for bitsandbytes.

            To resolve it, install PyTorch built for the ROCm version you want to use

            and set LD_LIBRARY_PATH to your ROCm install path, e.g.
            export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm-6.1.2/lib,
            """,
        )

        for pth in cudart_paths:
            print(f"* Found ROCm runtime at: {pth}")

def print_runtime_diagnostics() -> None:
    if HIP_ENVIRONMENT:
        _print_hip_runtime_diagnostics()
    else:
        _print_cuda_runtime_diagnostics()


        print("WARNING! ROCm runtime files not found in any environmental path.")
    elif len(cudart_paths) > 1:
        print_dedented(
            f"""
            Found duplicate ROCm runtime files (see below).

            We select the PyTorch default ROCm runtime, which is {torch.version.hip},
            but this might mismatch with the ROCm version that is needed for bitsandbytes.

            To resolve it, install PyTorch built for the ROCm version you want to use

            and set LD_LIBRARY_PATH to your ROCm install path, e.g.
            export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm-6.1.2/lib,
            """,
        )

        for pth in cudart_paths:
            print(f"* Found ROCm runtime at: {pth}")

def print_runtime_diagnostics() -> None:
    if HIP_ENVIRONMENT:
        _print_hip_runtime_diagnostics()
    else:
        _print_cuda_runtime_diagnostics()


bitsandbytes.diagnostics.main

import importlib
import platform
import sys
import traceback

import torch

from bitsandbytes import __version__ as bnb_version
from bitsandbytes.cextension import BNB_BACKEND
from bitsandbytes.consts import PACKAGE_GITHUB_URL
from bitsandbytes.cuda_specs import get_cuda_specs
from bitsandbytes.diagnostics.cuda import (
    print_diagnostics,
)
from bitsandbytes.diagnostics.utils import print_dedented, print_header

_RELATED_PACKAGES = [
    "accelerate",
    "diffusers",
    "numpy",
    "pip",
    "peft",
    "safetensors",
    "transformers",
    "triton",
    "trl",
]

def sanity_check():
    from bitsandbytes.optim import Adam

    p = torch.nn.Parameter(torch.rand(10, 10).cuda())
    a = torch.rand(10, 10).cuda()
    p1 = p.data.sum().item()
    adam = Adam([p])
    out = a * p
    loss = out.sum()
    loss.backward()
    adam.step()
    p2 = p.data.sum().item()
    assert p1 != p2

def get_package_version(name: str) -> str:
    try:
        version = importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        version = "not found"
    return version

def show_environment():
    """Simple utility to print out environment information."""

    print(f"Platform: {platform.platform()}")
    if platform.system() == "Linux":
        print(f"  libc: {'-'.join(platform.libc_ver())}")

    print(f"Python: {platform.python_version()}")

    print(f"PyTorch: {torch.__version__}")
    print(f"  CUDA: {torch.version.cuda or 'N/A'}")
    print(f"  HIP: {torch.version.hip or 'N/A'}")
    print(f"  XPU: {getattr(torch.version, 'xpu', 'N/A') or 'N/A'}")

    print("Related packages:")
    for pkg in _RELATED_PACKAGES:
        version = get_package_version(pkg)
        print(f"  {pkg}: {version}")

def main():
    print_header(f"bitsandbytes v{bnb_version}")
    show_environment()
    print_header("")

    cuda_specs = get_cuda_specs()

    if cuda_specs:
        print_diagnostics(cuda_specs)

    # TODO: There's a lot of noise in this; needs improvement.
    # print_cuda_runtime_diagnostics()

    if not torch.cuda.is_available():
        print(f"PyTorch says {BNB_BACKEND} is not available. Possible reasons:")
        print(f"1. {BNB_BACKEND} driver not installed")
        print("2. Using a CPU-only PyTorch build")
        print("3. No GPU detected")

    else:
        print(f"Checking that the library is importable and {BNB_BACKEND} is callable...")

        try:
            sanity_check()
            print("SUCCESS!")
            return
        except RuntimeError as e:
            if "not available in CPU-only" in str(e):
                print(
                    f"WARNING: {__package__} is currently running as CPU-only!\n"
                    "Therefore, 8-bit optimizers and GPU quantization are unavailable.\n\n"
                    f"If you think that this is so erroneously,\nplease report an issue!",
                )
            else:
                raise e
        except Exception:
            traceback.print_exc()

        print_dedented(
            f"""
            Above we output some debug information.
            Please provide this info when creating an issue via {PACKAGE_GITHUB_URL}/issues/new/choose
            WARNING: Please be sure to sanitize sensitive info from the output before posting it.
            """,
        )
        sys.exit(1)


from bitsandbytes.consts import PACKAGE_GITHUB_URL
from bitsandbytes.cuda_specs import get_cuda_specs
from bitsandbytes.diagnostics.cuda import (
    print_diagnostics,
)
from bitsandbytes.diagnostics.utils import print_dedented, print_header

_RELATED_PACKAGES = [
    "accelerate",
    "diffusers",
    "numpy",
    "pip",
    "peft",
    "safetensors",
    "transformers",
    "triton",
    "trl",
]

def sanity_check():
    from bitsandbytes.optim import Adam

    p = torch.nn.Parameter(torch.rand(10, 10).cuda())
    a = torch.rand(10, 10).cuda()
    p1 = p.data.sum().item()
    adam = Adam([p])
    out = a * p
    loss = out.sum()
    loss.backward()
    adam.step()
    p2 = p.data.sum().item()
    assert p1 != p2

def get_package_version(name: str) -> str:
    try:
        version = importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        version = "not found"
    return version

def show_environment():
    """Simple utility to print out environment information."""

    print(f"Platform: {platform.platform()}")
    if platform.system() == "Linux":
        print(f"  libc: {'-'.join(platform.libc_ver())}")

    print(f"Python: {platform.python_version()}")

    print(f"PyTorch: {torch.__version__}")


    "triton",
    "trl",
]

def sanity_check():
    from bitsandbytes.optim import Adam

    p = torch.nn.Parameter(torch.rand(10, 10).cuda())
    a = torch.rand(10, 10).cuda()
    p1 = p.data.sum().item()
    adam = Adam([p])
    out = a * p
    loss = out.sum()
    loss.backward()
    adam.step()
    p2 = p.data.sum().item()
    assert p1 != p2

def get_package_version(name: str) -> str:
    try:
        version = importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        version = "not found"
    return version

def show_environment():
    """Simple utility to print out environment information."""

    print(f"Platform: {platform.platform()}")
    if platform.system() == "Linux":
        print(f"  libc: {'-'.join(platform.libc_ver())}")

    print(f"Python: {platform.python_version()}")

    print(f"PyTorch: {torch.__version__}")
    print(f"  CUDA: {torch.version.cuda or 'N/A'}")
    print(f"  HIP: {torch.version.hip or 'N/A'}")
    print(f"  XPU: {getattr(torch.version, 'xpu', 'N/A') or 'N/A'}")

    print("Related packages:")
    for pkg in _RELATED_PACKAGES:
        version = get_package_version(pkg)
        print(f"  {pkg}: {version}")


    p = torch.nn.Parameter(torch.rand(10, 10).cuda())
    a = torch.rand(10, 10).cuda()
    p1 = p.data.sum().item()
    adam = Adam([p])
    out = a * p
    loss = out.sum()
    loss.backward()
    adam.step()
    p2 = p.data.sum().item()
    assert p1 != p2

def get_package_version(name: str) -> str:
    try:
        version = importlib.metadata.version(name)
    except importlib.metadata.PackageNotFoundError:
        version = "not found"
    return version

def show_environment():
    """Simple utility to print out environment information."""

    print(f"Platform: {platform.platform()}")
    if platform.system() == "Linux":
        print(f"  libc: {'-'.join(platform.libc_ver())}")

    print(f"Python: {platform.python_version()}")

    print(f"PyTorch: {torch.__version__}")
    print(f"  CUDA: {torch.version.cuda or 'N/A'}")
    print(f"  HIP: {torch.version.hip or 'N/A'}")
    print(f"  XPU: {getattr(torch.version, 'xpu', 'N/A') or 'N/A'}")

    print("Related packages:")
    for pkg in _RELATED_PACKAGES:
        version = get_package_version(pkg)
        print(f"  {pkg}: {version}")

def main():
    print_header(f"bitsandbytes v{bnb_version}")
    show_environment()
    print_header("")

    cuda_specs = get_cuda_specs()

    if cuda_specs:
        print_diagnostics(cuda_specs)

    # TODO: There's a lot of noise in this; needs improvement.
    # print_cuda_runtime_diagnostics()

    if not torch.cuda.is_available():
        print(f"PyTorch says {BNB_BACKEND} is not available. Possible reasons:")
        print(f"1. {BNB_BACKEND} driver not installed")
        print("2. Using a CPU-only PyTorch build")
        print("3. No GPU detected")


main

def show_environment():
    """Simple utility to print out environment information."""

    print(f"Platform: {platform.platform()}")
    if platform.system() == "Linux":
        print(f"  libc: {'-'.join(platform.libc_ver())}")

    print(f"Python: {platform.python_version()}")

    print(f"PyTorch: {torch.__version__}")
    print(f"  CUDA: {torch.version.cuda or 'N/A'}")
    print(f"  HIP: {torch.version.hip or 'N/A'}")
    print(f"  XPU: {getattr(torch.version, 'xpu', 'N/A') or 'N/A'}")

    print("Related packages:")
    for pkg in _RELATED_PACKAGES:
        version = get_package_version(pkg)
        print(f"  {pkg}: {version}")

def main():
    print_header(f"bitsandbytes v{bnb_version}")
    show_environment()
    print_header("")

    cuda_specs = get_cuda_specs()

    if cuda_specs:
        print_diagnostics(cuda_specs)

    # TODO: There's a lot of noise in this; needs improvement.
    # print_cuda_runtime_diagnostics()

    if not torch.cuda.is_available():
        print(f"PyTorch says {BNB_BACKEND} is not available. Possible reasons:")
        print(f"1. {BNB_BACKEND} driver not installed")
        print("2. Using a CPU-only PyTorch build")
        print("3. No GPU detected")

    else:
        print(f"Checking that the library is importable and {BNB_BACKEND} is callable...")

        try:
            sanity_check()
            print("SUCCESS!")
            return
        except RuntimeError as e:
            if "not available in CPU-only" in str(e):
                print(
                    f"WARNING: {__package__} is currently running as CPU-only!\n"
                    "Therefore, 8-bit optimizers and GPU quantization are unavailable.\n\n"
                    f"If you think that this is so erroneously,\nplease report an issue!",
                )
            else:
                raise e
        except Exception:
            traceback.print_exc()

        print_dedented(
            f"""
            Above we output some debug information.
            Please provide this info when creating an issue via {PACKAGE_GITHUB_URL}/issues/new/choose
            WARNING: Please be sure to sanitize sensitive info from the output before posting it.
            """,
        )
        sys.exit(1)


bitsandbytes.diagnostics.utils

import textwrap

HEADER_WIDTH = 60

def print_header(txt: str, width: int = HEADER_WIDTH, filler: str = "=") -> None:
    txt = f" {txt} " if txt else ""
    print(txt.center(width, filler))

def print_dedented(text):
    print("\n".join(textwrap.dedent(text).strip().split("\n")))


from functools import partial

import torch
import torch.nn as nn

from bitsandbytes.triton.dequantize_rowwise import dequantize_rowwise
from bitsandbytes.triton.int8_matmul_mixed_dequantize import (
    int8_matmul_mixed_dequantize,
)
from bitsandbytes.triton.int8_matmul_rowwise_dequantize import (
    int8_matmul_rowwise_dequantize,
)
from bitsandbytes.triton.quantize_columnwise_and_transpose import (
    quantize_columnwise_and_transpose,
)
from bitsandbytes.triton.quantize_global import (
    quantize_global,
    quantize_global_transpose,
)
from bitsandbytes.triton.quantize_rowwise import quantize_rowwise
from bitsandbytes.triton.triton_utils import is_triton_available

class _switchback_global(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))

        # rowwise quantize for X, global quantize for W
        X_int8, state_X = quantize_rowwise(X)
        W_int8, state_W = quantize_global(W)

        # save for backward.
        ctx.save_for_backward = X, W

        # matmult, fused dequant and add bias
        # call "mixed" because we are mixing rowwise quantized and global quantized
        return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        # reshape input to [N_out * L, D]
        G = G_3D.reshape(-1, G_3D.size(-1))

        grad_X = grad_W = grad_bias = None

        X, W = ctx.save_for_backward
        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, global quantize for W
            # for W, we also fuse the transpose operation because only A @ B^T is supported
            # so we transpose once then call .t() in the matmul
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_global_transpose(W)
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_vectorrize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))

        ctx.save_for_backward = X, W
        # rowwise quantize for X
        # columnwise quantize for W (first rowwise, transpose later)
        X_int8, state_X = quantize_rowwise(X)
        W_int8, state_W = quantize_rowwise(W)

        # matmult, fused dequant and add bias
        # call kernel which expects rowwise quantized X and W
        return int8_matmul_rowwise_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        X, W = ctx.save_for_backward

        G = G_3D.reshape(-1, G_3D.size(-1))

        grad_X = grad_W = grad_bias = None

        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, columnwise quantize for W and fused transpose
            # we call .t() for weight later because only A @ B^T is supported
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_columnwise_and_transpose(W)
            grad_X = int8_matmul_rowwise_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_global_mem_efficient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))
        X_3D_sz = X_3D.size()

        # rowwise quantize for X, global quantize for W
        X_int8, state_X = quantize_rowwise(X)
        del X
        W_int8, state_W = quantize_global(W)

        # save for backward.
        ctx.save_for_backward = X_int8, state_X, W_int8, state_W

        # matmult, fused dequant and add bias
        # call "mixed" because we are mixing rowwise quantized and global quantized
        return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D_sz[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        # reshape input to [N_out * L, D]
        G = G_3D.reshape(-1, G_3D.size(-1))
        G_3D_sz = G_3D.size()

        grad_X = grad_W = grad_bias = None

        X_int8, state_X, W_int8, state_W = ctx.save_for_backward
        if ctx.needs_input_grad[1]:
            real_X = dequantize_rowwise(X_int8, state_X)
            del X_int8
            grad_W = torch.matmul(G.t(), real_X.to(G.dtype))
            del real_X
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)
        if ctx.needs_input_grad[0]:
            G_int8, state_G = quantize_rowwise(G)
            del G
            W_int8 = W_int8.t().contiguous()
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(*G_3D_sz[:-1], -1)

        return grad_X, grad_W, grad_bias

class SwitchBackLinear(nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        device=None,
        dtype=None,
        vector_wise_quantization: bool = False,
        mem_efficient: bool = False,
    ):
        super().__init__(in_features, out_features, bias, device, dtype)

        if not is_triton_available():
            raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear.
                               Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""")

        # By default, we use the global quantization.
        self.vector_wise_quantization = vector_wise_quantization
        if self.vector_wise_quantization:
            self._fn = _switchback_vectorrize
            if mem_efficient:
                print("mem efficient is not supported for vector-wise quantization.")
                exit(1)
        else:
            if mem_efficient:
                self._fn = _switchback_global_mem_efficient
            else:
                self._fn = _switchback_global

    def prepare_for_eval(self):
        # If we just want to do eval, we can pre-quantize the weights instead of doing it on the forward pass.
        # Note this is experimental and not tested thoroughly.
        # Note this needs to be explicitly called with something like
        # def cond_prepare(m):
        #     if hasattr(m, "prepare_for_eval"):
        #         m.prepare_for_eval()
        # model.apply(cond_prepare)
        print("=> preparing for eval.")
        if self.vector_wise_quantization:
            W_int8, state_W = quantize_rowwise(self.weight)
        else:
            W_int8, state_W = quantize_global(self.weight)

        self.register_buffer("W_int8", W_int8)
        self.register_buffer("state_W", state_W)

        del self.weight

    def forward(self, x):
        if self.training:
            return self._fn.apply(x, self.weight, self.bias)
        else:
            # If it hasn't been "prepared for eval", run the standard forward pass.
            if not hasattr(self, "W_int8"):
                return self._fn.apply(x, self.weight, self.bias)

            # Otherwise, use pre-computed weights.
            X = x.view(-1, x.size(-1))
            X_int8, state_X = quantize_rowwise(X)

            if self.vector_wise_quantization:
                return int8_matmul_rowwise_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )
            else:
                return int8_matmul_mixed_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )

SwitchBackLinearGlobal = partial(SwitchBackLinear, vector_wise_quantization=False)
SwitchBackLinearGlobalMemEfficient = partial(SwitchBackLinear, vector_wise_quantization=False, mem_efficient=True)
SwitchBackLinearVectorwise = partial(SwitchBackLinear, vector_wise_quantization=True)

# This is just the standard linear function.
class StandardLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        X = input.view(-1, input.size(-1))

        ctx.save_for_backward(X, weight, bias)
        output = input.matmul(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output.view(*input.size()[:-1], -1)

    @staticmethod
    def backward(ctx, grad_output_3D):
        input, weight, bias = ctx.saved_tensors

        grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1))

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input.to(grad_output.dtype))
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

class StandardLinear(nn.Linear):
    def forward(self, x):
        return StandardLinearFunction.apply(x, self.weight, self.bias)


import torch.nn as nn

from bitsandbytes.triton.dequantize_rowwise import dequantize_rowwise
from bitsandbytes.triton.int8_matmul_mixed_dequantize import (
    int8_matmul_mixed_dequantize,
)
from bitsandbytes.triton.int8_matmul_rowwise_dequantize import (
    int8_matmul_rowwise_dequantize,
)
from bitsandbytes.triton.quantize_columnwise_and_transpose import (
    quantize_columnwise_and_transpose,
)
from bitsandbytes.triton.quantize_global import (
    quantize_global,
    quantize_global_transpose,
)
from bitsandbytes.triton.quantize_rowwise import quantize_rowwise
from bitsandbytes.triton.triton_utils import is_triton_available

class _switchback_global(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))

        # rowwise quantize for X, global quantize for W
        X_int8, state_X = quantize_rowwise(X)
        W_int8, state_W = quantize_global(W)

        # save for backward.
        ctx.save_for_backward = X, W

        # matmult, fused dequant and add bias
        # call "mixed" because we are mixing rowwise quantized and global quantized
        return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        # reshape input to [N_out * L, D]
        G = G_3D.reshape(-1, G_3D.size(-1))

        grad_X = grad_W = grad_bias = None

        X, W = ctx.save_for_backward
        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, global quantize for W
            # for W, we also fuse the transpose operation because only A @ B^T is supported
            # so we transpose once then call .t() in the matmul
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_global_transpose(W)
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_vectorrize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))

        ctx.save_for_backward = X, W
        # rowwise quantize for X
        # columnwise quantize for W (first rowwise, transpose later)
        X_int8, state_X = quantize_rowwise(X)
        W_int8, state_W = quantize_rowwise(W)

        # matmult, fused dequant and add bias
        # call kernel which expects rowwise quantized X and W
        return int8_matmul_rowwise_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):


        X, W = ctx.save_for_backward
        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, global quantize for W
            # for W, we also fuse the transpose operation because only A @ B^T is supported
            # so we transpose once then call .t() in the matmul
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_global_transpose(W)
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_vectorrize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))

        ctx.save_for_backward = X, W
        # rowwise quantize for X
        # columnwise quantize for W (first rowwise, transpose later)
        X_int8, state_X = quantize_rowwise(X)
        W_int8, state_W = quantize_rowwise(W)

        # matmult, fused dequant and add bias
        # call kernel which expects rowwise quantized X and W
        return int8_matmul_rowwise_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D.size()[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        X, W = ctx.save_for_backward

        G = G_3D.reshape(-1, G_3D.size(-1))

        grad_X = grad_W = grad_bias = None

        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, columnwise quantize for W and fused transpose
            # we call .t() for weight later because only A @ B^T is supported
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_columnwise_and_transpose(W)
            grad_X = int8_matmul_rowwise_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_global_mem_efficient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))
        X_3D_sz = X_3D.size()

        # rowwise quantize for X, global quantize for W
        X_int8, state_X = quantize_rowwise(X)
        del X
        W_int8, state_W = quantize_global(W)

        # save for backward.
        ctx.save_for_backward = X_int8, state_X, W_int8, state_W

        # matmult, fused dequant and add bias
        # call "mixed" because we are mixing rowwise quantized and global quantized
        return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D_sz[:-1], -1)


        grad_X = grad_W = grad_bias = None

        if ctx.needs_input_grad[0]:
            # rowwise quantize for G, columnwise quantize for W and fused transpose
            # we call .t() for weight later because only A @ B^T is supported
            G_int8, state_G = quantize_rowwise(G)
            W_int8, state_W = quantize_columnwise_and_transpose(W)
            grad_X = int8_matmul_rowwise_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(
                *G_3D.size()[:-1],
                -1,
            )
        if ctx.needs_input_grad[1]:
            # backward pass uses standard weight grad
            grad_W = torch.matmul(G.t(), X.to(G.dtype))
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)

        return grad_X, grad_W, grad_bias

class _switchback_global_mem_efficient(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X_3D, W, bias):
        # reshape input to [N * L, D]
        X = X_3D.view(-1, X_3D.size(-1))
        X_3D_sz = X_3D.size()

        # rowwise quantize for X, global quantize for W
        X_int8, state_X = quantize_rowwise(X)
        del X
        W_int8, state_W = quantize_global(W)

        # save for backward.
        ctx.save_for_backward = X_int8, state_X, W_int8, state_W

        # matmult, fused dequant and add bias
        # call "mixed" because we are mixing rowwise quantized and global quantized
        return int8_matmul_mixed_dequantize(X_int8, W_int8.t(), state_X, state_W, bias).view(*X_3D_sz[:-1], -1)

    @staticmethod
    def backward(ctx, G_3D):
        # reshape input to [N_out * L, D]
        G = G_3D.reshape(-1, G_3D.size(-1))
        G_3D_sz = G_3D.size()

        grad_X = grad_W = grad_bias = None

        X_int8, state_X, W_int8, state_W = ctx.save_for_backward
        if ctx.needs_input_grad[1]:
            real_X = dequantize_rowwise(X_int8, state_X)
            del X_int8
            grad_W = torch.matmul(G.t(), real_X.to(G.dtype))
            del real_X
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)
        if ctx.needs_input_grad[0]:
            G_int8, state_G = quantize_rowwise(G)
            del G
            W_int8 = W_int8.t().contiguous()
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(*G_3D_sz[:-1], -1)

        return grad_X, grad_W, grad_bias

class SwitchBackLinear(nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        device=None,
        dtype=None,
        vector_wise_quantization: bool = False,
        mem_efficient: bool = False,
    ):
        super().__init__(in_features, out_features, bias, device, dtype)

        if not is_triton_available():
            raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear.
                               Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""")

        # By default, we use the global quantization.


SwitchBackLinear

        grad_X = grad_W = grad_bias = None

        X_int8, state_X, W_int8, state_W = ctx.save_for_backward
        if ctx.needs_input_grad[1]:
            real_X = dequantize_rowwise(X_int8, state_X)
            del X_int8
            grad_W = torch.matmul(G.t(), real_X.to(G.dtype))
            del real_X
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)
        if ctx.needs_input_grad[0]:
            G_int8, state_G = quantize_rowwise(G)
            del G
            W_int8 = W_int8.t().contiguous()
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(*G_3D_sz[:-1], -1)

        return grad_X, grad_W, grad_bias

class SwitchBackLinear(nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        device=None,
        dtype=None,
        vector_wise_quantization: bool = False,
        mem_efficient: bool = False,
    ):
        super().__init__(in_features, out_features, bias, device, dtype)

        if not is_triton_available():
            raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear.
                               Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""")

        # By default, we use the global quantization.
        self.vector_wise_quantization = vector_wise_quantization
        if self.vector_wise_quantization:
            self._fn = _switchback_vectorrize
            if mem_efficient:
                print("mem efficient is not supported for vector-wise quantization.")
                exit(1)
        else:
            if mem_efficient:
                self._fn = _switchback_global_mem_efficient
            else:
                self._fn = _switchback_global

    def prepare_for_eval(self):
        # If we just want to do eval, we can pre-quantize the weights instead of doing it on the forward pass.
        # Note this is experimental and not tested thoroughly.
        # Note this needs to be explicitly called with something like
        # def cond_prepare(m):
        #     if hasattr(m, "prepare_for_eval"):
        #         m.prepare_for_eval()
        # model.apply(cond_prepare)
        print("=> preparing for eval.")
        if self.vector_wise_quantization:
            W_int8, state_W = quantize_rowwise(self.weight)
        else:
            W_int8, state_W = quantize_global(self.weight)

        self.register_buffer("W_int8", W_int8)
        self.register_buffer("state_W", state_W)

        del self.weight

    def forward(self, x):
        if self.training:
            return self._fn.apply(x, self.weight, self.bias)
        else:
            # If it hasn't been "prepared for eval", run the standard forward pass.
            if not hasattr(self, "W_int8"):
                return self._fn.apply(x, self.weight, self.bias)

            # Otherwise, use pre-computed weights.
            X = x.view(-1, x.size(-1))
            X_int8, state_X = quantize_rowwise(X)

            if self.vector_wise_quantization:
                return int8_matmul_rowwise_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )
            else:
                return int8_matmul_mixed_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )

SwitchBackLinearGlobal = partial(SwitchBackLinear, vector_wise_quantization=False)
SwitchBackLinearGlobalMemEfficient = partial(SwitchBackLinear, vector_wise_quantization=False, mem_efficient=True)
SwitchBackLinearVectorwise = partial(SwitchBackLinear, vector_wise_quantization=True)

# This is just the standard linear function.
class StandardLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        X = input.view(-1, input.size(-1))

        ctx.save_for_backward(X, weight, bias)
        output = input.matmul(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output.view(*input.size()[:-1], -1)

    @staticmethod


StandardLinearFunction

            X_int8, state_X = quantize_rowwise(X)

            if self.vector_wise_quantization:
                return int8_matmul_rowwise_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )
            else:
                return int8_matmul_mixed_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )

SwitchBackLinearGlobal = partial(SwitchBackLinear, vector_wise_quantization=False)
SwitchBackLinearGlobalMemEfficient = partial(SwitchBackLinear, vector_wise_quantization=False, mem_efficient=True)
SwitchBackLinearVectorwise = partial(SwitchBackLinear, vector_wise_quantization=True)

# This is just the standard linear function.
class StandardLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        X = input.view(-1, input.size(-1))

        ctx.save_for_backward(X, weight, bias)
        output = input.matmul(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output.view(*input.size()[:-1], -1)

    @staticmethod
    def backward(ctx, grad_output_3D):
        input, weight, bias = ctx.saved_tensors

        grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1))

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input.to(grad_output.dtype))
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

class StandardLinear(nn.Linear):
    def forward(self, x):
        return StandardLinearFunction.apply(x, self.weight, self.bias)


StandardLinear

        return output.view(*input.size()[:-1], -1)

    @staticmethod
    def backward(ctx, grad_output_3D):
        input, weight, bias = ctx.saved_tensors

        grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1))

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input.to(grad_output.dtype))
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

class StandardLinear(nn.Linear):
    def forward(self, x):
        return StandardLinearFunction.apply(x, self.weight, self.bias)


forward

    @staticmethod
    def backward(ctx, grad_output_3D):
        input, weight, bias = ctx.saved_tensors

        grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1))

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input.to(grad_output.dtype))
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

class StandardLinear(nn.Linear):
    def forward(self, x):
        return StandardLinearFunction.apply(x, self.weight, self.bias)


backward

SwitchBackLinearGlobal = partial(SwitchBackLinear, vector_wise_quantization=False)
SwitchBackLinearGlobalMemEfficient = partial(SwitchBackLinear, vector_wise_quantization=False, mem_efficient=True)
SwitchBackLinearVectorwise = partial(SwitchBackLinear, vector_wise_quantization=True)

# This is just the standard linear function.
class StandardLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        X = input.view(-1, input.size(-1))

        ctx.save_for_backward(X, weight, bias)
        output = input.matmul(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output.view(*input.size()[:-1], -1)

    @staticmethod
    def backward(ctx, grad_output_3D):
        input, weight, bias = ctx.saved_tensors

        grad_output = grad_output_3D.reshape(-1, grad_output_3D.size(-1))

        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.matmul(weight.to(grad_output.dtype)).view(*grad_output_3D.size()[:-1], -1)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().matmul(input.to(grad_output.dtype))
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

class StandardLinear(nn.Linear):
    def forward(self, x):
        return StandardLinearFunction.apply(x, self.weight, self.bias)


        grad_X = grad_W = grad_bias = None

        X_int8, state_X, W_int8, state_W = ctx.save_for_backward
        if ctx.needs_input_grad[1]:
            real_X = dequantize_rowwise(X_int8, state_X)
            del X_int8
            grad_W = torch.matmul(G.t(), real_X.to(G.dtype))
            del real_X
        if ctx.needs_input_grad[2]:
            grad_bias = G.sum(dim=0)
        if ctx.needs_input_grad[0]:
            G_int8, state_G = quantize_rowwise(G)
            del G
            W_int8 = W_int8.t().contiguous()
            grad_X = int8_matmul_mixed_dequantize(G_int8, W_int8.t(), state_G, state_W, None).view(*G_3D_sz[:-1], -1)

        return grad_X, grad_W, grad_bias

class SwitchBackLinear(nn.Linear):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        device=None,
        dtype=None,
        vector_wise_quantization: bool = False,
        mem_efficient: bool = False,
    ):
        super().__init__(in_features, out_features, bias, device, dtype)

        if not is_triton_available():
            raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear.
                               Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""")

        # By default, we use the global quantization.
        self.vector_wise_quantization = vector_wise_quantization
        if self.vector_wise_quantization:
            self._fn = _switchback_vectorrize
            if mem_efficient:
                print("mem efficient is not supported for vector-wise quantization.")
                exit(1)
        else:
            if mem_efficient:
                self._fn = _switchback_global_mem_efficient
            else:
                self._fn = _switchback_global

    def prepare_for_eval(self):
        # If we just want to do eval, we can pre-quantize the weights instead of doing it on the forward pass.
        # Note this is experimental and not tested thoroughly.
        # Note this needs to be explicitly called with something like
        # def cond_prepare(m):
        #     if hasattr(m, "prepare_for_eval"):
        #         m.prepare_for_eval()
        # model.apply(cond_prepare)
        print("=> preparing for eval.")
        if self.vector_wise_quantization:
            W_int8, state_W = quantize_rowwise(self.weight)
        else:
            W_int8, state_W = quantize_global(self.weight)

        self.register_buffer("W_int8", W_int8)
        self.register_buffer("state_W", state_W)

        del self.weight


    ):
        super().__init__(in_features, out_features, bias, device, dtype)

        if not is_triton_available():
            raise ImportError("""Could not import triton. Please install triton to use SwitchBackLinear.
                               Alternatively, you can use bnb.nn.SwitchBackLinearBnb, but it will be slower""")

        # By default, we use the global quantization.
        self.vector_wise_quantization = vector_wise_quantization
        if self.vector_wise_quantization:
            self._fn = _switchback_vectorrize
            if mem_efficient:
                print("mem efficient is not supported for vector-wise quantization.")
                exit(1)
        else:
            if mem_efficient:
                self._fn = _switchback_global_mem_efficient
            else:
                self._fn = _switchback_global

    def prepare_for_eval(self):
        # If we just want to do eval, we can pre-quantize the weights instead of doing it on the forward pass.
        # Note this is experimental and not tested thoroughly.
        # Note this needs to be explicitly called with something like
        # def cond_prepare(m):
        #     if hasattr(m, "prepare_for_eval"):
        #         m.prepare_for_eval()
        # model.apply(cond_prepare)
        print("=> preparing for eval.")
        if self.vector_wise_quantization:
            W_int8, state_W = quantize_rowwise(self.weight)
        else:
            W_int8, state_W = quantize_global(self.weight)

        self.register_buffer("W_int8", W_int8)
        self.register_buffer("state_W", state_W)

        del self.weight

    def forward(self, x):
        if self.training:
            return self._fn.apply(x, self.weight, self.bias)
        else:
            # If it hasn't been "prepared for eval", run the standard forward pass.
            if not hasattr(self, "W_int8"):
                return self._fn.apply(x, self.weight, self.bias)

            # Otherwise, use pre-computed weights.
            X = x.view(-1, x.size(-1))
            X_int8, state_X = quantize_rowwise(X)

            if self.vector_wise_quantization:
                return int8_matmul_rowwise_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(
                    *x.size()[:-1],
                    -1,
                )
            else:
                return int8_matmul_mixed_dequantize(X_int8, self.W_int8.t(), state_X, self.state_W, self.bias).view(


bitsandbytes.nn.parametrize

from functools import partial
from typing import Any, Literal, Optional

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P

from .. import functional as F

class Bnb4bitParametrization(nn.Module):
    """
    A parametrization module that handles dequantization of a 4-bit quantized parameter.

    The parameter data is expected to be already quantized when this parametrization is applied.
    This module will dequantize the parameter data to its original floating-point representation
    when the forward method is called (i.e. when the parameter is accessed).

    Args:
        quant_state (`F.QuantState`):
            The quantization state containing the necessary information for dequantization.
    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)

def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def replace_parameter_4bit(
    module: nn.Module,
    param_name: str,
    compress_statistics: bool = False,
    quant_type: Literal["nf4", "fp4"] = "nf4",
    blocksize: Optional[int] = None,
):
    """
    Replace a module parameter with a 4-bit quantized version using parametrization.

    This function quantizes an existing parameter in a PyTorch module to 4-bit precision
    and sets up parametrization to handle automatic dequantization during forward passes.
    The original parameter is replaced with quantized data, and a parametrization layer
    is registered to manage the quantization state and dequantization process.

    Additional, it registers a state dict post-hook to ensure that the quantization state
    is saved correctly when the model's state dict is saved.

    It is useful for MoE models or other scenarios where you want to quantize parameters
    outside of nn.Linear layers without changing the model's architecture.

    <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>

    Args:
        module (`nn.Module`):
            The PyTorch module containing the parameter to be quantized.
        param_name (`str`):
            The name of the parameter within the module to quantize.
        compress_statistics (`bool`, *optional*, defaults to `False`):
            Whether to compress quantization statistics to reduce memory usage.
        quant_type (`Literal["nf4", "fp4"]`, *optional*, defaults to `"nf4"`):
            The quantization format to use.
        blocksize (`int`, *optional*, defaults to `None`):
            The block size for quantization. If None, uses the default block size.

    Raises:
        AttributeError: If the module does not have the specified parameter.
        TypeError: If the specified attribute is not an instance of nn.Parameter.
    """

    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    # Quantize the original parameter.
    quantized_data, quant_state = F.quantize_4bit(
        original_param.data,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
    )

    # Replace the parameter with the quantized data.
    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    module.register_forward_pre_hook(_enable_parametrization_cache)
    module.register_forward_hook(_disable_parametrization_cache)

def _parametrized_state_dict_post_hook(
    module: nn.Module,
    state_dict: dict[str, Any],
    prefix: str,
    local_metadata: Any,
    *,
    param_name: str = "weight",
    **kwargs: dict[str, Any],
) -> None:
    """
    Hook to modify the state dict to include the quantization state.
    """

    original_key = f"{prefix}parametrizations.{param_name}.original"

    if original_key in state_dict:
        # Create a clean entry.
        # The `parametrizations.{param_name}.original` key will have the quantized data,
        # but we would like it to keep it in the state_dict as `{param_name}`.
        clean_key = f"{prefix}{param_name}"
        state_dict[clean_key] = state_dict.pop(original_key)

        assert P.is_parametrized(module, param_name)

        # Find the parametrization, which should have the quantization state.
        parametrization: Bnb4bitParametrization = next(
            filter(lambda x: isinstance(x, Bnb4bitParametrization), module.parametrizations[param_name]), None
        )

        assert parametrization is not None, "Parametrization not found for the parameter."

        quant_state = parametrization.quant_state

        # Next, we need to store the quantization state.
        if quant_state is not None:
            for k, v in quant_state.as_dict(packed=True).items():
                state_dict[f"{prefix}{param_name}.{k}"] = v


from functools import partial
from typing import Any, Literal, Optional

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P

from .. import functional as F

class Bnb4bitParametrization(nn.Module):
    """
    A parametrization module that handles dequantization of a 4-bit quantized parameter.

    The parameter data is expected to be already quantized when this parametrization is applied.
    This module will dequantize the parameter data to its original floating-point representation
    when the forward method is called (i.e. when the parameter is accessed).

    Args:
        quant_state (`F.QuantState`):
            The quantization state containing the necessary information for dequantization.
    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)

def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)


    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)

def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def replace_parameter_4bit(
    module: nn.Module,
    param_name: str,
    compress_statistics: bool = False,
    quant_type: Literal["nf4", "fp4"] = "nf4",
    blocksize: Optional[int] = None,
):
    """
    Replace a module parameter with a 4-bit quantized version using parametrization.

    This function quantizes an existing parameter in a PyTorch module to 4-bit precision
    and sets up parametrization to handle automatic dequantization during forward passes.
    The original parameter is replaced with quantized data, and a parametrization layer
    is registered to manage the quantization state and dequantization process.

    Additional, it registers a state dict post-hook to ensure that the quantization state
    is saved correctly when the model's state dict is saved.


def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def replace_parameter_4bit(
    module: nn.Module,
    param_name: str,
    compress_statistics: bool = False,
    quant_type: Literal["nf4", "fp4"] = "nf4",
    blocksize: Optional[int] = None,
):
    """
    Replace a module parameter with a 4-bit quantized version using parametrization.

    This function quantizes an existing parameter in a PyTorch module to 4-bit precision
    and sets up parametrization to handle automatic dequantization during forward passes.
    The original parameter is replaced with quantized data, and a parametrization layer
    is registered to manage the quantization state and dequantization process.

    Additional, it registers a state dict post-hook to ensure that the quantization state
    is saved correctly when the model's state dict is saved.

    It is useful for MoE models or other scenarios where you want to quantize parameters
    outside of nn.Linear layers without changing the model's architecture.

    <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>

    Args:
        module (`nn.Module`):
            The PyTorch module containing the parameter to be quantized.
        param_name (`str`):
            The name of the parameter within the module to quantize.
        compress_statistics (`bool`, *optional*, defaults to `False`):
            Whether to compress quantization statistics to reduce memory usage.
        quant_type (`Literal["nf4", "fp4"]`, *optional*, defaults to `"nf4"`):
            The quantization format to use.
        blocksize (`int`, *optional*, defaults to `None`):
            The block size for quantization. If None, uses the default block size.

    Raises:
        AttributeError: If the module does not have the specified parameter.
        TypeError: If the specified attribute is not an instance of nn.Parameter.
    """

    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    # Quantize the original parameter.
    quantized_data, quant_state = F.quantize_4bit(
        original_param.data,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
    )

    # Replace the parameter with the quantized data.
    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )


    # Quantize the original parameter.
    quantized_data, quant_state = F.quantize_4bit(
        original_param.data,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
    )

    # Replace the parameter with the quantized data.
    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    module.register_forward_pre_hook(_enable_parametrization_cache)


        quant_type=quant_type,
    )

    # Replace the parameter with the quantized data.
    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    module.register_forward_pre_hook(_enable_parametrization_cache)
    module.register_forward_hook(_disable_parametrization_cache)

def _parametrized_state_dict_post_hook(


    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)

def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    P._cache_enabled -= 1
    if not P._cache_enabled:
        P._cache = {}

def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    module.register_forward_pre_hook(_enable_parametrization_cache)
    module.register_forward_hook(_disable_parametrization_cache)

def _parametrized_state_dict_post_hook(
    module: nn.Module,
    state_dict: dict[str, Any],
    prefix: str,
    local_metadata: Any,
    *,
    param_name: str = "weight",
    **kwargs: dict[str, Any],
) -> None:
    """
    Hook to modify the state dict to include the quantization state.
    """

    original_key = f"{prefix}parametrizations.{param_name}.original"

    if original_key in state_dict:
        # Create a clean entry.
        # The `parametrizations.{param_name}.original` key will have the quantized data,


    P._cache_enabled += 1

def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    module.register_forward_pre_hook(_enable_parametrization_cache)
    module.register_forward_hook(_disable_parametrization_cache)

def _parametrized_state_dict_post_hook(
    module: nn.Module,
    state_dict: dict[str, Any],
    prefix: str,
    local_metadata: Any,
    *,
    param_name: str = "weight",
    **kwargs: dict[str, Any],
) -> None:
    """
    Hook to modify the state dict to include the quantization state.
    """

    original_key = f"{prefix}parametrizations.{param_name}.original"

    if original_key in state_dict:
        # Create a clean entry.
        # The `parametrizations.{param_name}.original` key will have the quantized data,
        # but we would like it to keep it in the state_dict as `{param_name}`.
        clean_key = f"{prefix}{param_name}"
        state_dict[clean_key] = state_dict.pop(original_key)

        assert P.is_parametrized(module, param_name)

        # Find the parametrization, which should have the quantization state.
        parametrization: Bnb4bitParametrization = next(
            filter(lambda x: isinstance(x, Bnb4bitParametrization), module.parametrizations[param_name]), None
        )

        assert parametrization is not None, "Parametrization not found for the parameter."

        quant_state = parametrization.quant_state

        # Next, we need to store the quantization state.
        if quant_state is not None:
            for k, v in quant_state.as_dict(packed=True).items():
                state_dict[f"{prefix}{param_name}.{k}"] = v


import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P

from .. import functional as F

class Bnb4bitParametrization(nn.Module):
    """
    A parametrization module that handles dequantization of a 4-bit quantized parameter.

    The parameter data is expected to be already quantized when this parametrization is applied.
    This module will dequantize the parameter data to its original floating-point representation
    when the forward method is called (i.e. when the parameter is accessed).

    Args:
        quant_state (`F.QuantState`):
            The quantization state containing the necessary information for dequantization.
    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)

def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")


forward

class Bnb4bitParametrization(nn.Module):
    """
    A parametrization module that handles dequantization of a 4-bit quantized parameter.

    The parameter data is expected to be already quantized when this parametrization is applied.
    This module will dequantize the parameter data to its original floating-point representation
    when the forward method is called (i.e. when the parameter is accessed).

    Args:
        quant_state (`F.QuantState`):
            The quantization state containing the necessary information for dequantization.
    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)

def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)