← back to Exo

.typings/mlx_lm/models/gemma4_text.pyi

180 lines

from dataclasses import dataclass
from typing import Any, Dict, List, Optional

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

from .base import BaseModelArgs
from .cache import KVCache, RotatingKVCache
from .switch_layers import SwitchGLU

@dataclass
class ModelArgs(BaseModelArgs):
    model_type: str
    hidden_size: int
    num_hidden_layers: int
    intermediate_size: int
    num_attention_heads: int
    head_dim: int
    global_head_dim: int
    global_partial_rotary_factor: float
    rms_norm_eps: float
    vocab_size: int
    vocab_size_per_layer_input: int
    num_key_value_heads: int
    num_global_key_value_heads: Optional[int]
    num_kv_shared_layers: int
    pad_token_id: int
    hidden_size_per_layer_input: int
    rope_traditional: bool
    partial_rotary_factor: float
    rope_parameters: Optional[Dict[str, Any]]
    sliding_window: int
    sliding_window_pattern: int
    max_position_embeddings: int
    attention_k_eq_v: bool
    final_logit_softcapping: float
    use_double_wide_mlp: bool
    enable_moe_block: bool
    num_experts: Optional[int]
    top_k_experts: Optional[int]
    moe_intermediate_size: Optional[int]
    layer_types: Optional[List[str]]
    tie_word_embeddings: bool

    def __post_init__(self) -> None: ...

class MLP(nn.Module):
    gate_proj: nn.Linear
    down_proj: nn.Linear
    up_proj: nn.Linear

    def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...

class Router(nn.Module):
    proj: nn.Linear
    scale: mx.array
    per_expert_scale: mx.array

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...

class Experts(nn.Module):
    switch_glu: SwitchGLU

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(
        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
    ) -> mx.array: ...

class Attention(nn.Module):
    layer_idx: int
    layer_type: str
    is_sliding: bool
    head_dim: int
    n_heads: int
    n_kv_heads: int
    use_k_eq_v: bool
    scale: float
    q_proj: nn.Linear
    k_proj: nn.Linear
    v_proj: nn.Linear
    o_proj: nn.Linear
    q_norm: nn.Module
    k_norm: nn.Module
    v_norm: nn.Module
    rope: nn.Module

    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...

class DecoderLayer(nn.Module):
    layer_idx: int
    layer_type: str
    self_attn: Attention
    mlp: MLP
    enable_moe: bool
    router: Router
    experts: Experts
    input_layernorm: nn.Module
    post_attention_layernorm: nn.Module
    pre_feedforward_layernorm: nn.Module
    post_feedforward_layernorm: nn.Module
    post_feedforward_layernorm_1: nn.Module
    post_feedforward_layernorm_2: nn.Module
    pre_feedforward_layernorm_2: nn.Module
    hidden_size_per_layer_input: int
    per_layer_input_gate: Optional[nn.Linear]
    per_layer_projection: Optional[nn.Linear]
    post_per_layer_input_norm: Optional[nn.Module]
    layer_scalar: mx.array

    def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = ...,
        cache: Optional[Any] = ...,
        per_layer_input: Optional[mx.array] = ...,
        shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
        offset: Optional[mx.array] = ...,
    ) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...

class Gemma4TextModel(nn.Module):
    config: ModelArgs
    vocab_size: int
    window_size: int
    sliding_window_pattern: int
    num_hidden_layers: int
    embed_tokens: nn.Embedding
    embed_scale: float
    layers: list[DecoderLayer]
    norm: nn.Module
    hidden_size_per_layer_input: int
    embed_tokens_per_layer: Optional[nn.Embedding]
    per_layer_model_projection: Optional[nn.Linear]
    per_layer_projection_norm: Optional[nn.Module]
    previous_kvs: list[int]

    def __init__(self, config: ModelArgs) -> None: ...
    def __call__(
        self,
        inputs: Optional[mx.array] = ...,
        cache: Optional[list[Any]] = ...,
        input_embeddings: Optional[mx.array] = ...,
        per_layer_inputs: Optional[mx.array] = ...,
    ) -> mx.array: ...
    def _get_per_layer_inputs(
        self,
        input_ids: Optional[mx.array],
        input_embeddings: Optional[mx.array] = ...,
    ) -> mx.array: ...
    def _project_per_layer_inputs(
        self,
        input_embeddings: mx.array,
        per_layer_inputs: Optional[mx.array] = ...,
    ) -> mx.array: ...
    def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...

class Model(nn.Module):
    args: ModelArgs
    model_type: str
    model: Gemma4TextModel
    final_logit_softcapping: float
    tie_word_embeddings: bool
    lm_head: nn.Linear

    def __init__(self, args: ModelArgs) -> None: ...
    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
    def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
    @property
    def layers(self) -> list[DecoderLayer]: ...
    @property
    def head_dim(self) -> int: ...
    @property
    def n_kv_heads(self) -> int: ...
    @property
    def quant_predicate(self) -> Any: ...
    def make_cache(self) -> list[KVCache | RotatingKVCache]: ...