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

from functools import partial

import mlx.core as mx
import mlx.nn as nn

class QuantizedSwitchLinear(nn.Module):
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        num_experts: int,
        bias: bool = ...,
        group_size: int = ...,
        bits: int = ...,
        mode: str = ...,
    ) -> None: ...
    @property
    def input_dims(self):  # -> int:
        ...
    @property
    def output_dims(self):  # -> int:
        ...
    @property
    def num_experts(self):  # -> int:
        ...
    def __call__(self, x, indices, sorted_indices=...):  # -> array:
        ...

class SwitchLinear(nn.Module):
    def __init__(
        self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ...
    ) -> None: ...
    @property
    def input_dims(self):  # -> int:
        ...
    @property
    def output_dims(self):  # -> int:
        ...
    @property
    def num_experts(self):  # -> int:
        ...
    def __call__(self, x, indices, sorted_indices=...): ...
    def to_quantized(
        self, group_size: int = ..., bits: int = ..., mode: str = ...
    ):  # -> QuantizedSwitchLinear:
        ...

@partial(mx.compile, shapeless=True)
def swiglu(x, gate): ...

class SwiGLU(nn.Module):
    def __init__(self) -> None: ...
    def __call__(self, x, gate): ...

class SwitchGLU(nn.Module):
    gate_proj: SwitchLinear
    up_proj: SwitchLinear
    down_proj: SwitchLinear
    activation: SwiGLU

    def __init__(
        self,
        input_dims: int,
        hidden_dims: int,
        num_experts: int,
        activation=...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x, indices) -> mx.array: ...

class SwitchMLP(nn.Module):
    fc1: SwitchLinear
    fc2: SwitchLinear

    def __init__(
        self,
        input_dims: int,
        hidden_dims: int,
        num_experts: int,
        activation=...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x, indices) -> mx.array: ...
