"""
This type stub file was generated by pyright.
"""

import mlx.nn as nn

class LoRALinear(nn.Module):
    @staticmethod
    def from_base(
        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
    ):  # -> LoRALinear:
        ...
    def fuse(self, dequantize: bool = ...):  # -> QuantizedLinear | Linear:
        ...
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        r: int = ...,
        dropout: float = ...,
        scale: float = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x):  # -> array:
        ...

class LoRASwitchLinear(nn.Module):
    @staticmethod
    def from_base(
        linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
    ):  # -> LoRASwitchLinear:
        ...
    def fuse(self, dequantize: bool = ...):  # -> QuantizedSwitchLinear | SwitchLinear:
        ...
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        num_experts: int,
        r: int = ...,
        dropout: float = ...,
        scale: float = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x, indices, sorted_indices=...): ...

class LoRAEmbedding(nn.Module):
    @staticmethod
    def from_base(
        embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
    ):  # -> LoRAEmbedding:
        ...
    def fuse(self, dequantize: bool = ...):  # -> QuantizedEmbedding | Embedding:
        ...
    def __init__(
        self,
        num_embeddings: int,
        dims: int,
        r: int = ...,
        dropout: float = ...,
        scale: float = ...,
    ) -> None: ...
    def __call__(self, x):  # -> array:
        ...
    def as_linear(self, x):  # -> array:
        ...
