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

import mlx.nn as nn

class DoRALinear(nn.Module):
    @staticmethod
    def from_base(
        linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
    ):  # -> DoRALinear:
        ...
    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 set_linear(self, linear):  # -> None:
        """
        Set the self.linear layer and recompute self.m.
        """
        ...

    def __call__(self, x): ...

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