← back to Exo

.typings/mflux/models/common/weights/mapping/weight_mapping.pyi

24 lines

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

import mlx.core as mx
from dataclasses import dataclass
from typing import Callable, List, Optional, Protocol

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

@dataclass
class WeightTarget:
    to_pattern: str
    from_pattern: List[str]
    transform: Optional[Callable[[mx.array], mx.array]] = ...
    required: bool = ...
    max_blocks: Optional[int] = ...

class WeightMapping(Protocol):
    @staticmethod
    def get_mapping() -> List[WeightTarget]: ...