Use 2GB buffer for more accurate bandwidth measurement

- Increase buffer size from 512MB to 2GB to better saturate memory bus
- Use 2D array shape to avoid issues with very large 1D arrays
- Improves accuracy from ~75% to ~82% of theoretical peak

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.cuda_typings/openai_harmony/__init__.pyi

@@ -1,20 +0,0 @@
-from enum import Enum
-
-class HarmonyEncodingName(Enum):
-    HARMONY_GPT_OSS = ...
-
-class HarmonyEncoding: ...
-class HarmonyError(Exception): ...
-
-class Role(Enum):
-    ASSISTANT = ...
-
-class StreamableParser:
-    last_content_delta: str
-    current_channel: str | None
-    current_recipient: str | None
-
-    def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
-    def process(self, token_id: int) -> None: ...
-
-def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...

.cuda_typings/pynvml/__init__.pyi

@@ -1,17 +0,0 @@
-class NvmlMemoryInfo:
-    used: int
-    total: int
-    free: int
-
-class NvmlUtilizationRates:
-    gpu: int
-    memory: int
-
-def nvmlInit() -> None: ...
-def nvmlShutdown() -> None: ...
-def nvmlDeviceGetCount() -> int: ...
-def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
-def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
-def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
-def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
-def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...

.cuda_typings/torch/__init__.pyi

@@ -1,61 +0,0 @@
-from typing import Any, Sequence
-
-from torch import backends as backends
-from torch import cuda as cuda
-from torch import distributed as distributed
-
-__version__: str
-
-class version:
-    cuda: str
-
-class dtype: ...
-
-bfloat16: dtype
-float16: dtype
-float32: dtype
-int8: dtype
-int32: dtype
-int64: dtype
-long: dtype
-float8_e4m3fn: dtype
-
-class Tensor:
-    shape: Sequence[int]
-    dtype: dtype
-    def __getitem__(self, key: Any) -> Tensor: ...
-    def __setitem__(self, key: Any, value: Any) -> None: ...
-    def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
-    def cpu(self) -> Tensor: ...
-    def detach(self) -> Tensor: ...
-    def clone(self) -> Tensor: ...
-    def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
-    def view(self, *shape: Any) -> Tensor: ...
-    def squeeze(self, dim: int = ...) -> Tensor: ...
-    def unsqueeze(self, dim: int) -> Tensor: ...
-    def permute(self, *dims: int) -> Tensor: ...
-    def float(self) -> Tensor: ...
-    def numpy(self) -> Any: ...
-    def numel(self) -> int: ...
-    def nelement(self) -> int: ...
-    @property
-    def is_cuda(self) -> bool: ...
-    @property
-    def device(self) -> device: ...
-    def __len__(self) -> int: ...
-    def data_ptr(self) -> int: ...
-    def tolist(self) -> Any: ...
-    def abs(self) -> Tensor: ...
-    def max(self) -> Tensor: ...
-    def mean(self) -> Tensor: ...
-    def sum(self, dim: int = ...) -> Tensor: ...
-    def item(self) -> float: ...
-
-def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
-def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
-def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
-def from_numpy(ndarray: Any) -> Tensor: ...
-def inference_mode() -> Any: ...
-
-class device:
-    def __init__(self, type: str, index: int = ...) -> None: ...

.cuda_typings/torch/backends/__init__.pyi

@@ -1 +0,0 @@
-from torch.backends import cuda as cuda

.cuda_typings/torch/backends/cuda/__init__.pyi

@@ -1 +0,0 @@
-def is_built() -> bool: ...

.cuda_typings/torch/cuda/__init__.pyi

@@ -1,10 +0,0 @@
-class _DeviceProperties:
-    total_memory: int
-
-def is_available() -> bool: ...
-def get_device_name(device: int) -> str: ...
-def get_device_properties(device: int) -> _DeviceProperties: ...
-def empty_cache() -> None: ...
-def mem_get_info() -> tuple[int, int]: ...
-def synchronize() -> None: ...
-def max_memory_allocated() -> int: ...

.cuda_typings/torch/distributed/__init__.pyi

@@ -1,2 +0,0 @@
-def is_initialized() -> bool: ...
-def destroy_process_group() -> None: ...

.cuda_typings/vllm/__init__.pyi

@@ -1 +0,0 @@
-__version__: str

.cuda_typings/vllm/config.pyi

@@ -1,2 +0,0 @@
-class ModelConfig:
-    max_model_len: int

.cuda_typings/vllm/engine/arg_utils.pyi

@@ -1,18 +0,0 @@
-from dataclasses import dataclass
-
-@dataclass
-class EngineArgs:
-    model: str = ...
-    served_model_name: str | list[str] | None = ...
-    tokenizer: str | None = ...
-    trust_remote_code: bool = ...
-    dtype: str = ...
-    seed: int = ...
-    max_model_len: int | None = ...
-    gpu_memory_utilization: float = ...
-    enforce_eager: bool = ...
-    tensor_parallel_size: int = ...
-    pipeline_parallel_size: int = ...
-    quantization: str | None = ...
-    load_format: str = ...
-    enable_sleep_mode: bool = ...

.cuda_typings/vllm/outputs.pyi

@@ -1,17 +0,0 @@
-class CompletionOutput:
-    index: int
-    text: str
-    token_ids: list[int]
-    cumulative_logprob: float | None
-    logprobs: object | None
-    finish_reason: str | None
-    stop_reason: int | str | None
-
-    def finished(self) -> bool: ...
-
-class RequestOutput:
-    request_id: str
-    prompt: str | None
-    prompt_token_ids: list[int] | None
-    outputs: list[CompletionOutput]
-    finished: bool

.cuda_typings/vllm/sampling_params.pyi

@@ -1,11 +0,0 @@
-class SamplingParams:
-    n: int
-    temperature: float
-    top_p: float
-    top_k: int
-    min_p: float
-    seed: int | None
-    stop: str | list[str] | None
-    max_tokens: int | None
-    logprobs: int | None
-    repetition_penalty: float

.cuda_typings/vllm/tokenizers/__init__.pyi

@@ -1,3 +0,0 @@
-from vllm.tokenizers.protocol import TokenizerLike
-
-__all__ = ["TokenizerLike"]

.cuda_typings/vllm/tokenizers/protocol.pyi

@@ -1,15 +0,0 @@
-from typing import Protocol
-
-class TokenizerLike(Protocol):
-    @property
-    def eos_token_id(self) -> int: ...
-    @property
-    def vocab_size(self) -> int: ...
-    def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
-    def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
-    def apply_chat_template(
-        self,
-        messages: list[dict[str, str]],
-        tools: list[dict[str, object]] | None = ...,
-        **kwargs: object,
-    ) -> str | list[int]: ...

.cuda_typings/vllm/v1/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/core/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/core/kv_cache_manager.pyi

@@ -1,24 +0,0 @@
-from collections.abc import Sequence
-
-from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
-from vllm.v1.kv_cache_interface import KVCacheConfig
-
-class KVCacheBlocks:
-    blocks: tuple[Sequence[KVCacheBlock], ...]
-    def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
-    def get_block_ids(self) -> tuple[list[int], ...]: ...
-
-class KVCacheManager:
-    block_pool: BlockPool
-    kv_cache_config: KVCacheConfig
-    enable_caching: bool
-    num_kv_cache_groups: int
-    coordinator: object
-    def __init__(self, *args: object, **kwargs: object) -> None: ...
-    def allocate_slots(
-        self, request: object, num_new_tokens: int, *args: object, **kwargs: object
-    ) -> KVCacheBlocks | None: ...
-    def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
-    def create_kv_cache_blocks(
-        self, blocks: tuple[list[KVCacheBlock], ...]
-    ) -> KVCacheBlocks: ...

.cuda_typings/vllm/v1/core/kv_cache_utils.pyi

@@ -1,16 +0,0 @@
-class KVCacheBlock:
-    block_id: int
-    ref_cnt: int
-    def __init__(self, block_id: int) -> None: ...
-
-class FreeKVCacheBlockQueue:
-    def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
-    def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
-
-class BlockPool:
-    blocks: list[KVCacheBlock]
-    free_block_queue: FreeKVCacheBlockQueue
-    num_gpu_blocks: int
-    enable_caching: bool
-    def get_num_free_blocks(self) -> int: ...
-    def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...

.cuda_typings/vllm/v1/engine/llm_engine.pyi

@@ -1,22 +0,0 @@
-from vllm.config import ModelConfig
-from vllm.engine.arg_utils import EngineArgs
-from vllm.outputs import RequestOutput
-from vllm.sampling_params import SamplingParams
-from vllm.tokenizers import TokenizerLike
-
-class LLMEngine:
-    tokenizer: TokenizerLike | None
-    model_config: ModelConfig
-
-    @classmethod
-    def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
-    def add_request(
-        self,
-        request_id: str,
-        prompt: str,
-        params: SamplingParams,
-        arrival_time: float | None = ...,
-    ) -> None: ...
-    def step(self) -> list[RequestOutput]: ...
-    def has_unfinished_requests(self) -> bool: ...
-    def get_tokenizer(self) -> TokenizerLike: ...

.cuda_typings/vllm/v1/kv_cache_interface.pyi

@@ -1,23 +0,0 @@
-from dataclasses import dataclass
-
-@dataclass
-class KVCacheSpec:
-    block_size: int
-    num_kv_heads: int
-    head_size: int
-
-@dataclass
-class KVCacheGroupSpec:
-    layer_names: list[str]
-    kv_cache_spec: KVCacheSpec
-
-@dataclass
-class KVCacheTensorSpec:
-    shared_by: list[str]
-    size: int
-
-@dataclass
-class KVCacheConfig:
-    num_blocks: int
-    kv_cache_groups: list[KVCacheGroupSpec]
-    kv_cache_tensors: list[KVCacheTensorSpec]

.cuda_typings/vllm/v1/request.pyi

@@ -1,6 +0,0 @@
-class Request:
-    request_id: str
-    prompt_token_ids: list[int] | None
-    num_prompt_tokens: int
-    num_computed_tokens: int
-    num_tokens: int

.cuda_typings/vllm/v1/worker/__init__.pyi

@@ -1 +0,0 @@
-

.cuda_typings/vllm/v1/worker/gpu_model_runner.pyi

@@ -1,24 +0,0 @@
-import torch
-
-class _CompilationConfig:
-    static_forward_context: dict[str, object]
-
-class _ModelConfig:
-    hf_config: object
-
-class GPUModelRunner:
-    kv_caches: list[torch.Tensor]
-    compilation_config: _CompilationConfig
-    model_config: _ModelConfig | None
-    def _allocate_kv_cache_tensors(
-        self, kv_cache_config: object
-    ) -> dict[str, torch.Tensor]: ...
-    def initialize_kv_cache_tensors(
-        self, kv_cache_config: object, kernel_block_sizes: list[int]
-    ) -> dict[str, torch.Tensor]: ...
-    def _reshape_kv_cache_tensors(
-        self,
-        kv_cache_config: object,
-        raw_tensors: dict[str, torch.Tensor],
-        kernel_block_sizes: list[int],
-    ) -> dict[str, torch.Tensor]: ...

.cuda_typings/vllm/v1/worker/gpu_worker.pyi

@@ -1,6 +0,0 @@
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-class Worker:
-    model_runner: GPUModelRunner
-    def determine_available_memory(self) -> int: ...
-    def initialize_from_config(self, kv_cache_config: object) -> None: ...

.cuda_typings/vllm/v1/worker/utils.pyi

@@ -1 +0,0 @@
-def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...

.github/actions/typecheck/action.yml

@@ -0,0 +1,12 @@
+name: Type Check
+
+description: "Run type checker"
+
+runs:
+  using: "composite"
+  steps:
+    - name: Run type checker
+      run: |
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
+        nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
+      shell: bash

.github/workflows/build-app.yml

@@ -1,16 +1,5 @@
 name: Build EXO macOS DMG
 
-# Release workflow:
-# 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown
-# 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0
-# 3. This workflow builds, signs, and notarizes the DMG
-# 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown)
-# 5. DMG and appcast.xml are uploaded to S3
-# 6. The draft GitHub Release is published with the DMG attached
-#
-# For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional.
-# If no draft exists, a release is auto-created with generated notes.
-
 on:
   workflow_dispatch:
   push:
@@ -22,10 +11,8 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
-    permissions:
-      contents: write
     env:
-      SPARKLE_VERSION: 2.9.0-beta.1
+      SPARKLE_VERSION: 2.8.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
       SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }}
       SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }}
@@ -100,52 +87,6 @@ jobs:
             exit 1
           fi
 
-      - name: Fetch and validate release notes
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          # Find draft release by name using gh release list (more reliable with default token)
-          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
-          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
-
-          if [[ -z "$DRAFT_EXISTS" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: No draft release found for tag $GITHUB_REF_NAME"
-            echo "Please create a draft release with release notes before pushing the tag."
-            exit 1
-          fi
-
-          # Fetch full release details via API to get body and ID
-          echo "Found draft release, fetching details..."
-          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
-
-          # Extract release notes
-          NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
-          if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "Draft release has no notes (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: Draft release exists but has no release notes"
-            echo "Please add release notes to the draft release before pushing the tag."
-            exit 1
-          fi
-
-          # Save release ID for later publishing
-          RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id')
-          echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV
-          echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV
-
-          echo "Found draft release (ID: $RELEASE_ID), saving release notes..."
-          echo "$NOTES" > /tmp/release_notes.md
-          echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV
-
       # ============================================================
       # Install dependencies
       # ============================================================
@@ -363,28 +304,6 @@ jobs:
             $CHANNEL_FLAG \
             .
 
-      - name: Inject release notes into appcast
-        if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true'
-        env:
-          RELEASE_VERSION: ${{ env.RELEASE_VERSION }}
-        run: |
-          # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+)
-          export NOTES=$(cat "$RELEASE_NOTES_FILE")
-
-          # Insert description after the enclosure tag for this version
-          awk '
-            /<enclosure[^>]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) {
-              print
-              print "            <description sparkle:format=\"markdown\"><![CDATA["
-              print ENVIRON["NOTES"]
-              print "            ]]></description>"
-              next
-            }
-            { print }
-          ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml
-
-          echo "Injected markdown release notes for version $RELEASE_VERSION"
-
       # ============================================================
       # Upload artifacts
       # ============================================================
@@ -396,7 +315,7 @@ jobs:
           path: output/EXO-${{ env.RELEASE_VERSION }}.dmg
 
       - name: Upload to S3
-        if: env.SPARKLE_S3_BUCKET != ''
+        if: env.SPARKLE_S3_BUCKET != '' && github.ref_type == 'tag'
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -412,37 +331,8 @@ jobs:
             PREFIX="${PREFIX}/"
           fi
           DMG_NAME="EXO-${RELEASE_VERSION}.dmg"
-
-          if [[ "${{ github.ref_type }}" != "tag" ]]; then
-            aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-${GITHUB_SHA}.dmg"
-            exit 0
-          fi
-
           aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}${DMG_NAME}"
           if [[ "$IS_ALPHA" != "true" ]]; then
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
             aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache
           fi
-
-      - name: Publish GitHub Release
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg"
-
-          if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then
-            # Update the draft release with the tag and upload DMG
-            gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \
-              -f tag_name="$GITHUB_REF_NAME" \
-              -F draft=false
-            gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber
-            echo "Published release $GITHUB_REF_NAME with DMG attached"
-          else
-            # Alpha without draft release - create one with auto-generated notes
-            gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \
-              --title "$GITHUB_REF_NAME" \
-              --generate-notes \
-              --prerelease
-            echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes"
-          fi

.github/workflows/pipeline.yml

@@ -8,6 +8,92 @@ on:
       - main
 
 jobs:
+  typecheck:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          lfs: false
+
+      - uses: cachix/install-nix-action@v31
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+
+      - uses: cachix/cachix-action@v14
+        name: Configure Cachix
+        with:
+          name: exo
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+
+      - name: Configure git user
+        run: |
+          git config --local user.email "github-actions@users.noreply.github.com"
+          git config --local user.name  "github-actions bot"
+        shell: bash
+
+      - name: Pull LFS files
+        run: |
+          echo "Pulling Git LFS files..."
+          git lfs pull
+        shell: bash
+
+      - name: Setup Nix Environment
+        run: |
+          echo "Checking for nix installation..."
+          
+          # Check if nix binary exists directly
+          if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
+            echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
+            export PATH="/nix/var/nix/profiles/default/bin:$PATH"
+            echo "PATH=$PATH" >> $GITHUB_ENV
+            nix --version
+          elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
+            echo "Found nix profile script, sourcing..."
+            source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
+            nix --version
+          elif command -v nix >/dev/null 2>&1; then
+            echo "Nix already in PATH"
+            nix --version
+          else
+            echo "Nix not found. Debugging info:"
+            echo "Contents of /nix/var/nix/profiles/default/:"
+            ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
+            echo "Contents of /nix/var/nix/profiles/default/bin/:"
+            ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
+            exit 1
+          fi
+        shell: bash
+
+      - name: Configure basedpyright include for local MLX
+        run: |
+          RUNNER_LABELS='${{ toJSON(runner.labels) }}'
+          if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
+            if [ -d "/Users/Shared/mlx" ]; then
+              echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
+              awk '
+                BEGIN { in=0 }
+                /^\[tool\.basedpyright\]/ { in=1; print; next }
+                in && /^\[/ { in=0 }  # next section
+                in && /^[ \t]*include[ \t]*=/ {
+                  print "include = [\"/Users/Shared/mlx\"]"
+                  next
+                }
+                { print }
+              ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
+
+              echo "New [tool.basedpyright] section:"
+              sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
+            else
+              echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
+            fi
+          else
+            echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
+          fi
+        shell: bash
+
+      - uses: ./.github/actions/typecheck
+
   nix:
     name: Build and check (${{ matrix.system }})
     runs-on: ${{ matrix.runner }}
@@ -37,63 +123,6 @@ jobs:
           name: exo
           authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
 
-      - name: Build Metal packages (macOS only)
-        if: runner.os == 'macOS'
-        run: |
-          # Try to build metal-toolchain first (may succeed via cachix cache hit)
-          if nix build .#metal-toolchain 2>/dev/null; then
-            echo "metal-toolchain built successfully (likely cache hit)"
-          else
-            echo "metal-toolchain build failed, extracting from Xcode..."
-
-            NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw="
-            NAR_NAME="metal-toolchain-17C48.nar"
-
-            # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS
-            WORK_DIR="${RUNNER_TEMP}/metal-work"
-            mkdir -p "$WORK_DIR"
-
-            # Download the Metal toolchain component
-            xcodebuild -downloadComponent MetalToolchain
-
-            # Find and mount the DMG
-            DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1)
-            if [ -z "$DMG_PATH" ]; then
-              echo "Error: Could not find Metal toolchain DMG"
-              exit 1
-            fi
-
-            echo "Found DMG at: $DMG_PATH"
-            hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg"
-
-            # Copy the toolchain
-            cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export"
-            hdiutil detach "${WORK_DIR}/metal-dmg"
-
-            # Create NAR and add to store
-            nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}"
-            STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}")
-            echo "Added NAR to store: $STORE_PATH"
-
-            # Verify the hash matches
-            ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}")
-            if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then
-              echo "Warning: NAR hash mismatch!"
-              echo "Expected: $NAR_HASH"
-              echo "Actual:   $ACTUAL_HASH"
-              echo "The metal-toolchain.nix may need updating"
-            fi
-
-            # Clean up
-            rm -rf "$WORK_DIR"
-
-            # Retry the build now that NAR is in store
-            nix build .#metal-toolchain
-          fi
-
-          # Build mlx (depends on metal-toolchain)
-          nix build .#mlx
-
       - name: Build all Nix outputs
         run: |
           nix flake show --json | jq -r '
@@ -105,16 +134,3 @@ jobs:
 
       - name: Run nix flake check
         run: nix flake check
-
-      - name: Run pytest (macOS only)
-        if: runner.os == 'macOS'
-        run: |
-          # Build the test environment (requires relaxed sandbox for uv2nix on macOS)
-          TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths)
-
-          # Run pytest outside sandbox (needs GPU access for MLX)
-          export HOME="$RUNNER_TEMP"
-          export EXO_TESTS=1
-          export EXO_DASHBOARD_DIR="$PWD/dashboard/" 
-          export EXO_RESOURCES_DIR="$PWD/resources" 
-          $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib

.gitignore

@@ -28,13 +28,3 @@ target/
 dashboard/build/
 dashboard/node_modules/
 dashboard/.svelte-kit/
-
-# host config snapshots
-hosts_*.json
-.swp
-
-# bench files
-bench/**/*.json
-
-# tmp
-tmp/models

.mlx_typings/mlx/core/__init__.pyi

@@ -215,22 +215,6 @@ class StreamContext:
         traceback: object | None = ...,
     ) -> None: ...
 
-def device_info() -> dict[str, str | int]:
-    """
-    Get information about the GPU device and system settings.
-
-    Currently returns:
-
-    * ``architecture``
-    * ``max_buffer_size``
-    * ``max_recommended_working_set_size``
-    * ``memory_size``
-    * ``resource_limit``
-
-    Returns:
-        dict: A dictionary with string keys and string or integer values.
-    """
-
 def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
     """
     Element-wise absolute value.
@@ -1155,7 +1139,7 @@ class array:
     ) -> array:
         """See :func:`flatten`."""
 
-    def reshape(self, *shape: int, stream: Stream | Device | None = ...) -> array:
+    def reshape(self, *shape, stream: Stream | Device | None = ...) -> array:
         """
         Equivalent to :func:`reshape` but the shape can be passed either as a
         :obj:`tuple` or as separate arguments.
@@ -1238,7 +1222,7 @@ class array:
     ) -> array:
         """See :func:`swapaxes`."""
 
-    def transpose(self, *axes: int, stream: Stream | Device | None = ...) -> array:
+    def transpose(self, *axes, stream: Stream | Device | None = ...) -> array:
         """
         Equivalent to :func:`transpose` but the axes can be passed either as
         a tuple or as separate arguments.
@@ -2382,7 +2366,7 @@ class custom_function:
 def default_device() -> Device:
     """Get the default device."""
 
-def default_stream(device: Device | DeviceType) -> Stream:
+def default_stream(device: Device) -> Stream:
     """Get the device's default stream."""
 
 def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:

.mlx_typings/mlx/nn/__init__.pyi

@@ -2,7 +2,8 @@
 This type stub file was generated by pyright.
 """
 
+from layers import *
+from utils import *
+
 from . import init as init
 from . import losses as losses
-from .layers import *
-from .utils import *

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -0,0 +1,20 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from activations import *
+from base import *
+from containers import *
+from convolution import *
+from convolution_transpose import *
+from distributed import *
+from dropout import *
+from embedding import *
+from linear import *
+from normalization import *
+from pooling import *
+from positional_encoding import *
+from quantized import *
+from recurrent import *
+from transformer import *
+from upsample import *

.mlx_typings/mlx/nn/layers/activations.pyi

@@ -6,7 +6,7 @@ from functools import partial
 from typing import Any
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 @partial(mx.compile, shapeless=True)
 def sigmoid(x: mx.array) -> mx.array:

.mlx_typings/mlx/nn/layers/base.pyi

@@ -200,7 +200,7 @@ class Module(dict):
     ) -> mx.MX_ARRAY_TREE:  # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
         """Return the submodules that do not contain other modules."""
 
-    def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module:
+    def update(self, parameters: dict, strict: bool = ...) -> Module:
         """Replace the parameters of this Module with the provided ones in the
         dict of dicts and lists.
 

.mlx_typings/mlx/nn/layers/containers.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Callable
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class Sequential(Module):
     """A layer that calls the passed callables in order.

.mlx_typings/mlx/nn/layers/convolution.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class Conv1d(Module):
     """Applies a 1-dimensional convolution over the multi-channel input sequence.
@@ -30,10 +30,6 @@ class Conv1d(Module):
         bias (bool, optional): If ``True`` add a learnable bias to the output.
             Default: ``True``
     """
-
-    weight: mx.array
-    bias: mx.array | None
-    groups: int
     def __init__(
         self,
         in_channels: int,

.mlx_typings/mlx/nn/layers/convolution_transpose.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class ConvTranspose1d(Module):
     """Applies a 1-dimensional transposed convolution over the multi-channel input sequence.

.mlx_typings/mlx/nn/layers/distributed.pyi

@@ -6,7 +6,7 @@ from functools import lru_cache
 from typing import Callable, Optional, Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 from mlx.nn.layers.linear import Linear
 
 @lru_cache

.mlx_typings/mlx/nn/layers/dropout.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class Dropout(Module):
     r"""Randomly zero a portion of the elements during training.

.mlx_typings/mlx/nn/layers/embedding.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 from .quantized import QuantizedEmbedding
 

.mlx_typings/mlx/nn/layers/linear.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Any
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 from .quantized import QuantizedLinear
 
@@ -40,10 +40,6 @@ class Linear(Module):
         bias (bool, optional): If set to ``False`` then the layer will
           not use a bias. Default is ``True``.
     """
-
-    weight: mx.array
-    bias: mx.array | None
-
     def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
     def __call__(self, x: mx.array) -> mx.array: ...
     def to_quantized(

.mlx_typings/mlx/nn/layers/normalization.pyi

@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
 """
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class InstanceNorm(Module):
     r"""Applies instance normalization [1] on the inputs.
@@ -88,9 +88,6 @@ class RMSNorm(Module):
         dims (int): The feature dimension of the input to normalize over
         eps (float): A small additive constant for numerical stability
     """
-
-    weight: mx.array
-
     def __init__(self, dims: int, eps: float = ...) -> None: ...
     def __call__(self, x) -> mx.array: ...
 

.mlx_typings/mlx/nn/layers/pooling.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Optional, Tuple, Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class _Pool(Module):
     def __init__(

.mlx_typings/mlx/nn/layers/positional_encoding.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Optional
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class RoPE(Module):
     """Implements the rotary positional encoding.

.mlx_typings/mlx/nn/layers/quantized.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Callable, Optional, Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 def quantize(
     model: Module,

.mlx_typings/mlx/nn/layers/recurrent.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Callable, Optional
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class RNN(Module):
     r"""An Elman recurrent layer.

.mlx_typings/mlx/nn/layers/transformer.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Any, Callable, Optional
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 class MultiHeadAttention(Module):
     """Implements the scaled dot product attention with multiple heads.

.mlx_typings/mlx/nn/layers/upsample.pyi

@@ -5,7 +5,7 @@ This type stub file was generated by pyright.
 from typing import Literal, Tuple, Union
 
 import mlx.core as mx
-from .base import Module
+from base import Module
 
 def upsample_nearest(x: mx.array, scale_factor: Tuple) -> mx.array: ...
 def upsample_linear(

.mlx_typings/mlx/utils.pyi

@@ -2,15 +2,12 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from mlx.core import MX_ARRAY_TREE
 
 def tree_map(
-    fn: Callable[..., Any],
-    tree: Any,
-    *rest: Any,
-    is_leaf: Callable[..., bool] | None = ...,
+    fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
 ) -> Any:
     """Applies ``fn`` to the leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -47,11 +44,11 @@ def tree_map(
     """
 
 def tree_map_with_path(
-    fn: Callable[..., Any],
+    fn: Callable,
     tree: Any,
     *rest: Any,
-    is_leaf: Callable[..., bool] | None = ...,
-    path: str | None = ...,
+    is_leaf: Optional[Callable] = ...,
+    path: Optional[Any] = ...,
 ) -> Any:
     """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
     returns a new collection with the results.
@@ -83,9 +80,9 @@ def tree_map_with_path(
 def tree_flatten(
     tree: Any,
     prefix: str = ...,
-    is_leaf: Callable[..., bool] | None = ...,
-    destination: list[tuple[str, Any]] | dict[str, Any] | None = ...,
-) -> list[tuple[str, Any]] | dict[str, Any]:
+    is_leaf: Optional[Callable] = ...,
+    destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ...,
+) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
     """Flattens a Python tree to a list of key, value tuples.
 
     The keys are using the dot notation to define trees of arbitrary depth and
@@ -121,7 +118,7 @@ def tree_flatten(
             the Python tree.
     """
 
-def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any:
+def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
     """Recreate a Python tree from its flat representation.
 
     .. code-block:: python

.mlx_typings/mlx_lm/generate.pyi

@@ -73,11 +73,9 @@ class GenerationResponse:
     finish_reason: Optional[str] = ...
 
 def maybe_quantize_kv_cache(
-    prompt_cache: Any,
-    quantized_kv_start: int | None,
-    kv_group_size: int | None,
-    kv_bits: int | None,
-) -> None: ...
+    prompt_cache, quantized_kv_start, kv_group_size, kv_bits
+):  # -> None:
+    ...
 def generate_step(
     prompt: mx.array,
     model: nn.Module,
@@ -254,14 +252,7 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
-    caches: Optional[List[List[Any]]]
 
-def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
-def _make_cache(
-    model: Any, left_padding: Any, max_kv_size: Optional[int]
-) -> List[Any]: ...
-def _merge_caches(caches: Any) -> List[Any]: ...
 @dataclass
 class Batch:
     uids: List[int]
@@ -270,71 +261,39 @@ class Batch:
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
-    tokens: List[mx.array]
-    def __len__(self) -> int: ...
-    def filter(self, keep_idx: List[int]) -> None: ...
-    def extend(self, other: "Batch") -> None: ...
-    def extract_cache(self, idx: int) -> List[Any]: ...
+    def __len__(self):  # -> int:
+        ...
+    def filter(self, keep_idx: List[int]):  # -> None:
+        ...
+    def extend(self, other):  # -> None:
+        ...
 
 class BatchGenerator:
-    model: Any
-    max_kv_size: Optional[int]
-    prefill_step_size: int
-    unprocessed_prompts: List[Any]
-    active_batch: Optional[Batch]
-    prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
-    _stats: BatchStats
-
     @dataclass
     class Response:
         uid: int
         token: int
         logprobs: mx.array
         finish_reason: Optional[str]
-        prompt_cache: Any
 
     def __init__(
         self,
-        model: Any,
+        model,
         max_tokens: int = ...,
-        stop_tokens: Optional[set[int]] = ...,
+        stop_tokens: Optional[set] = ...,
         sampler: Optional[Callable[[mx.array], mx.array]] = ...,
-        logits_processors: Optional[
-            List[Callable[[mx.array, mx.array], mx.array]]
-        ] = ...,
         completion_batch_size: int = ...,
         prefill_batch_size: int = ...,
         prefill_step_size: int = ...,
-        prompt_progress_callback: Optional[
-            Callable[[List[Tuple[int, int, int]]], None]
-        ] = ...,
-        max_kv_size: Optional[int] = ...,
     ) -> None: ...
-    def close(self) -> None: ...
     def insert(
-        self,
-        prompts: Any,
-        max_tokens: Union[List[int], int, None] = ...,
-        caches: Any = ...,
-        samplers: Optional[List[Any]] = ...,
-        logits_processors: Optional[List[Any]] = ...,
-    ) -> List[int]: ...
-    def remove(
-        self, uids: List[int], return_prompt_caches: bool = ...
-    ) -> Optional[dict[int, List[Any]]]: ...
-    def stats(self) -> BatchStats: ...
-    def next(self) -> List[Response]: ...
-    def _process_prompts(self, prompts: List[Any]) -> Batch: ...
-    def _step(
-        self,
-        input_tokens: mx.array,
-        prompt_cache: List[Any],
-        samplers: Optional[List[Any]],
-        logits_processors: Optional[List[Any]],
-        tokens: List[mx.array],
-    ) -> Tuple[mx.array, List[mx.array]]: ...
+        self, prompts, max_tokens: Union[List[int], int, None] = ...
+    ):  # -> list[Any]:
+        ...
+    def stats(self):  # -> BatchStats:
+        ...
+    def next(self):  # -> list[Any]:
+        ...
 
 def batch_generate(
     model,

.mlx_typings/mlx_lm/models/cache.pyi

@@ -2,21 +2,18 @@
 This type stub file was generated by pyright.
 """
 
-from typing import Any, Dict, List, Literal, Optional, Protocol, Self
+from typing import Any, Dict, List, Optional, Protocol, Literal, Self
 
-import mlx.core as mx
 import mlx.nn as nn
 from mlx.core import array
+import mlx.core as mx
 
 class Cache(Protocol):
     keys: mx.array
     values: mx.array
-    offset: int
-    def update_and_fetch(
-        self, keys: mx.array, values: mx.array
-    ) -> tuple[mx.array, mx.array]: ...
+    def update_and_fetch(self, keys: mx.array, values: mx.array) -> None: ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
 
@@ -90,16 +87,14 @@ def create_attention_mask(
 class _BaseCache(Cache):
     keys: mx.array
     values: mx.array
-    offset: int
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self) -> tuple[mx.array, mx.array]: ...
     @state.setter
     def state(self, v) -> None: ...
     @property
     def meta_state(self) -> Literal[""]: ...
     @meta_state.setter
     def meta_state(self, v) -> None: ...
-    def trim(self, n: int) -> int: ...
     def is_trimmable(self) -> Literal[False]: ...
     @classmethod
     def from_state(cls, state, meta_state) -> Self: ...
@@ -115,13 +110,15 @@ class ConcatenateKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> tuple[Any | array, Any | array]:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -131,7 +128,10 @@ class QuantizedKVCache(_BaseCache):
     def update_and_fetch(self, keys, values):  # -> Any:
         ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(
+        self,
+    ):  # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -143,7 +143,8 @@ class QuantizedKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
         ...
 
@@ -155,30 +156,22 @@ class KVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ) -> tuple[array, array]: ...
     @state.setter
     def state(self, v) -> None: ...
     def is_trimmable(self):  # -> Literal[True]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
-    def make_mask(
-        self, *args: Any, **kwargs: Any
-    ) -> mx.array | Literal["causal"] | None: ...
+    def make_mask(self, *args, **kwargs):  # -> array | Literal['causal'] | None:
+        ...
 
 class RotatingKVCache(_BaseCache):
     step = ...
-    keys: mx.array | None
-    values: mx.array | None
-    keep: int
-    max_size: int
-    _idx: int
     def __init__(self, max_size, keep=...) -> None: ...
-    def _trim(
-        self, trim_size: int, v: mx.array, append: mx.array | None = ...
-    ) -> mx.array: ...
     def update_and_fetch(
         self, keys, values
     ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
@@ -186,7 +179,8 @@ class RotatingKVCache(_BaseCache):
     @property
     def state(
         self,
-    ) -> tuple[mx.array | None, mx.array | None]: ...
+    ):  # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -198,7 +192,8 @@ class RotatingKVCache(_BaseCache):
         ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...
     ) -> QuantizedKVCache: ...
@@ -213,7 +208,8 @@ class ArraysCache(_BaseCache):
         ...
     def __getitem__(self, idx): ...
     @property
-    def state(self) -> tuple[mx.array | None, mx.array | None]: ...
+    def state(self):  # -> list[Any | array] | list[array]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...
@@ -227,7 +223,8 @@ class ArraysCache(_BaseCache):
         In-place extend this cache with the other cache.
         """
 
-    def make_mask(self, N: int) -> mx.array | None: ...
+    def make_mask(self, N: int):  # -> array | None:
+        ...
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -238,7 +235,8 @@ class ChunkedKVCache(KVCache):
         ...
     def update_and_fetch(self, keys, values):  # -> tuple[array, array]:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n):  # -> int:
+        ...
     @property
     def meta_state(self):  # -> tuple[str, ...]:
         ...
@@ -251,9 +249,10 @@ class CacheList(_BaseCache):
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...
-    def trim(self, n: int) -> int: ...
+    def trim(self, n): ...
     @property
-    def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
+    def state(self):  # -> list[Any]:
+        ...
     @state.setter
     def state(self, v):  # -> None:
         ...

.mlx_typings/mlx_lm/models/deepseek_v3.pyi

@@ -5,7 +5,6 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx_lm.models.mla import MultiLinear
 
 from .base import BaseModelArgs
 from .switch_layers import SwitchGLU
@@ -61,10 +60,7 @@ class DeepseekV3Attention(nn.Module):
     q_b_proj: nn.Linear
     kv_a_proj_with_mqa: nn.Linear
     kv_a_layernorm: nn.RMSNorm
-    # kv_b_proj: nn.Linear
-    embed_q: MultiLinear
-    unembed_out: MultiLinear
-
+    kv_b_proj: nn.Linear
     o_proj: nn.Linear
     rope: Any
 

.mlx_typings/mlx_lm/models/switch_layers.pyi

@@ -73,9 +73,6 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array: ...
 
 class SwitchMLP(nn.Module):
-    fc1: SwitchLinear
-    fc2: SwitchLinear
-
     def __init__(
         self,
         input_dims: int,

.mlx_typings/mlx_lm/sample_utils.pyi

@@ -48,7 +48,7 @@ def make_logits_processors(
     logit_bias: Optional[Dict[int, float]] = ...,
     repetition_penalty: Optional[float] = ...,
     repetition_context_size: Optional[int] = ...,
-) -> list[Callable[[mx.array, mx.array], mx.array]]:
+):  # -> list[Any]:
     """
     Make logits processors for use with ``generate_step``.
 

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -2,6 +2,7 @@
 This type stub file was generated by pyright.
 """
 
+from functools import partial
 from pathlib import Path
 from typing import Any
 
@@ -38,11 +39,11 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    def reset(self) -> None: ...
-    def add_token(self, token: int) -> None: ...
-    def finalize(self) -> None: ...
+    def reset(self): ...
+    def add_token(self, token): ...
+    def finalize(self): ...
     @property
-    def last_segment(self) -> str:
+    def last_segment(self):
         """Return the last segment of readable text since last time this property was accessed."""
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
@@ -107,21 +108,16 @@ class TokenizerWrapper:
     _tokenizer: PreTrainedTokenizerFast
     eos_token_id: int | None
     eos_token: str | None
-    eos_token_ids: list[int] | set[int] | None
     bos_token_id: int | None
     bos_token: str | None
     vocab_size: int
     all_special_tokens: list[str]
-    think_start: str | None
-    think_end: str | None
-    think_start_id: int | None
-    think_end_id: int | None
 
     def __init__(
         self,
         tokenizer: Any,
         detokenizer_class: Any = ...,
-        eos_token_ids: list[int] | set[int] | None = ...,
+        eos_token_ids: list[int] | None = ...,
         chat_template: Any = ...,
         tool_parser: Any = ...,
         tool_call_start: str | None = ...,

.prettierrc

@@ -0,0 +1,3 @@
+{
+  "useTabs": true
+}

.typings/mflux/__init__.pyi

@@ -1,7 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import os
-
-if "TOKENIZERS_PARALLELISM" not in os.environ: ...

.typings/mflux/callbacks/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/callbacks/callback.pyi

@@ -1,48 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import Protocol
-
-import mlx.core as mx
-import PIL.Image
-import tqdm
-from mflux.models.common.config.config import Config
-
-class BeforeLoopCallback(Protocol):
-    def call_before_loop(
-        self,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        canny_image: PIL.Image.Image | None = ...,
-        depth_image: PIL.Image.Image | None = ...,
-    ) -> None: ...
-
-class InLoopCallback(Protocol):
-    def call_in_loop(
-        self,
-        t: int,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        time_steps: tqdm,
-    ) -> None: ...
-
-class AfterLoopCallback(Protocol):
-    def call_after_loop(
-        self, seed: int, prompt: str, latents: mx.array, config: Config
-    ) -> None: ...
-
-class InterruptCallback(Protocol):
-    def call_interrupt(
-        self,
-        t: int,
-        seed: int,
-        prompt: str,
-        latents: mx.array,
-        config: Config,
-        time_steps: tqdm,
-    ) -> None: ...

.typings/mflux/callbacks/callback_registry.pyi

@@ -1,25 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import TYPE_CHECKING
-
-from mflux.callbacks.callback import (
-    AfterLoopCallback,
-    BeforeLoopCallback,
-    InLoopCallback,
-    InterruptCallback,
-)
-from mflux.callbacks.generation_context import GenerationContext
-from mflux.models.common.config.config import Config
-
-if TYPE_CHECKING: ...
-
-class CallbackRegistry:
-    def __init__(self) -> None: ...
-    def register(self, callback) -> None: ...
-    def start(self, seed: int, prompt: str, config: Config) -> GenerationContext: ...
-    def before_loop_callbacks(self) -> list[BeforeLoopCallback]: ...
-    def in_loop_callbacks(self) -> list[InLoopCallback]: ...
-    def after_loop_callbacks(self) -> list[AfterLoopCallback]: ...
-    def interrupt_callbacks(self) -> list[InterruptCallback]: ...

.typings/mflux/callbacks/generation_context.pyi

@@ -1,30 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import TYPE_CHECKING
-
-import mlx.core as mx
-import PIL.Image
-import tqdm
-from mflux.callbacks.callback_registry import CallbackRegistry
-from mflux.models.common.config.config import Config
-
-if TYPE_CHECKING: ...
-
-class GenerationContext:
-    def __init__(
-        self, registry: CallbackRegistry, seed: int, prompt: str, config: Config
-    ) -> None: ...
-    def before_loop(
-        self,
-        latents: mx.array,
-        *,
-        canny_image: PIL.Image.Image | None = ...,
-        depth_image: PIL.Image.Image | None = ...,
-    ) -> None: ...
-    def in_loop(self, t: int, latents: mx.array, time_steps: tqdm = ...) -> None: ...
-    def after_loop(self, latents: mx.array) -> None: ...
-    def interruption(
-        self, t: int, latents: mx.array, time_steps: tqdm = ...
-    ) -> None: ...

.typings/mflux/cli/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/cli/defaults/defaults.pyi

@@ -1,22 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import os
-
-BATTERY_PERCENTAGE_STOP_LIMIT = ...
-CONTROLNET_STRENGTH = ...
-DEFAULT_DEV_FILL_GUIDANCE = ...
-DEFAULT_DEPTH_GUIDANCE = ...
-DIMENSION_STEP_PIXELS = ...
-GUIDANCE_SCALE = ...
-GUIDANCE_SCALE_KONTEXT = ...
-IMAGE_STRENGTH = ...
-MODEL_CHOICES = ...
-MODEL_INFERENCE_STEPS = ...
-QUANTIZE_CHOICES = ...
-if os.environ.get("MFLUX_CACHE_DIR"):
-    MFLUX_CACHE_DIR = ...
-else:
-    MFLUX_CACHE_DIR = ...
-MFLUX_LORA_CACHE_DIR = ...

.typings/mflux/models/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/cli/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/config/__init__.pyi

@@ -1,8 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.config.config import Config
-from mflux.models.common.config.model_config import ModelConfig
-
-__all__ = ["Config", "ModelConfig"]

.typings/mflux/models/common/config/config.pyi

@@ -1,67 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-from typing import Any
-
-import mlx.core as mx
-from mflux.models.common.config.model_config import ModelConfig
-from tqdm import tqdm
-
-logger = ...
-
-class Config:
-    def __init__(
-        self,
-        model_config: ModelConfig,
-        num_inference_steps: int = ...,
-        height: int = ...,
-        width: int = ...,
-        guidance: float = ...,
-        image_path: Path | str | None = ...,
-        image_strength: float | None = ...,
-        depth_image_path: Path | str | None = ...,
-        redux_image_paths: list[Path | str] | None = ...,
-        redux_image_strengths: list[float] | None = ...,
-        masked_image_path: Path | str | None = ...,
-        controlnet_strength: float | None = ...,
-        scheduler: str = ...,
-    ) -> None: ...
-    @property
-    def height(self) -> int: ...
-    @property
-    def width(self) -> int: ...
-    @width.setter
-    def width(self, value):  # -> None:
-        ...
-    @property
-    def image_seq_len(self) -> int: ...
-    @property
-    def guidance(self) -> float: ...
-    @property
-    def num_inference_steps(self) -> int: ...
-    @property
-    def precision(self) -> mx.Dtype: ...
-    @property
-    def num_train_steps(self) -> int: ...
-    @property
-    def image_path(self) -> Path | None: ...
-    @property
-    def image_strength(self) -> float | None: ...
-    @property
-    def depth_image_path(self) -> Path | None: ...
-    @property
-    def redux_image_paths(self) -> list[Path] | None: ...
-    @property
-    def redux_image_strengths(self) -> list[float] | None: ...
-    @property
-    def masked_image_path(self) -> Path | None: ...
-    @property
-    def init_time_step(self) -> int: ...
-    @property
-    def time_steps(self) -> tqdm: ...
-    @property
-    def controlnet_strength(self) -> float | None: ...
-    @property
-    def scheduler(self) -> Any: ...

.typings/mflux/models/common/config/model_config.pyi

@@ -1,87 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from functools import lru_cache
-from typing import Literal
-
-import mlx.core as mx
-
-class ModelConfig:
-    precision: mx.Dtype = ...
-    def __init__(
-        self,
-        priority: int,
-        aliases: list[str],
-        model_name: str,
-        base_model: str | None,
-        controlnet_model: str | None,
-        custom_transformer_model: str | None,
-        num_train_steps: int | None,
-        max_sequence_length: int | None,
-        supports_guidance: bool | None,
-        requires_sigma_shift: bool | None,
-        transformer_overrides: dict | None = ...,
-    ) -> None: ...
-    @staticmethod
-    @lru_cache
-    def dev() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def schnell() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_kontext() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_fill() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_redux() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_depth() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_controlnet_canny() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def schnell_controlnet_canny() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_controlnet_upscaler() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def dev_fill_catvton() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def krea_dev() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def flux2_klein_4b() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def flux2_klein_9b() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def qwen_image() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def qwen_image_edit() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def fibo() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def z_image_turbo() -> ModelConfig: ...
-    @staticmethod
-    @lru_cache
-    def seedvr2_3b() -> ModelConfig: ...
-    def x_embedder_input_dim(self) -> int: ...
-    def is_canny(self) -> bool: ...
-    @staticmethod
-    def from_name(
-        model_name: str, base_model: Literal["dev", "schnell", "krea-dev"] | None = ...
-    ) -> ModelConfig: ...
-
-AVAILABLE_MODELS = ...

.typings/mflux/models/common/latent_creator/__init__.pyi

@@ -1,7 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/latent_creator/latent_creator.pyi

@@ -1,50 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-from typing import TYPE_CHECKING, TypeAlias
-
-import mlx.core as mx
-from mflux.models.common.vae.tiling_config import TilingConfig
-from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
-from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
-from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator
-from mflux.models.z_image.latent_creator.z_image_latent_creator import (
-    ZImageLatentCreator,
-)
-from mlx import nn
-
-if TYPE_CHECKING:
-    LatentCreatorType: TypeAlias = type[
-        FiboLatentCreator | FluxLatentCreator | QwenLatentCreator | ZImageLatentCreator
-    ]
-
-class Img2Img:
-    def __init__(
-        self,
-        vae: nn.Module,
-        latent_creator: LatentCreatorType,
-        sigmas: mx.array,
-        init_time_step: int,
-        image_path: str | Path | None,
-        tiling_config: TilingConfig | None = ...,
-    ) -> None: ...
-
-class LatentCreator:
-    @staticmethod
-    def create_for_txt2img_or_img2img(
-        seed: int, height: int, width: int, img2img: Img2Img
-    ) -> mx.array: ...
-    @staticmethod
-    def encode_image(
-        vae: nn.Module,
-        image_path: str | Path,
-        height: int,
-        width: int,
-        tiling_config: TilingConfig | None = ...,
-    ) -> mx.array: ...
-    @staticmethod
-    def add_noise_by_interpolation(
-        clean: mx.array, noise: mx.array, sigma: float
-    ) -> mx.array: ...

.typings/mflux/models/common/lora/__init__.pyi

@@ -1,3 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""

.typings/mflux/models/common/lora/layer/fused_linear_lora_layer.pyi

@@ -1,13 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear
-from mlx import nn
-
-class FusedLoRALinear(nn.Module):
-    def __init__(
-        self, base_linear: nn.Linear | nn.QuantizedLinear, loras: list[LoRALinear]
-    ) -> None: ...
-    def __call__(self, x):  # -> array:
-        ...

.typings/mflux/models/common/lora/layer/linear_lora_layer.pyi

@@ -1,22 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mlx import nn
-
-class LoRALinear(nn.Module):
-    @staticmethod
-    def from_linear(
-        linear: nn.Linear | nn.QuantizedLinear, r: int = ..., scale: float = ...
-    ):  # -> LoRALinear:
-        ...
-    def __init__(
-        self,
-        input_dims: int,
-        output_dims: int,
-        r: int = ...,
-        scale: float = ...,
-        bias: bool = ...,
-    ) -> None: ...
-    def __call__(self, x):  # -> array:
-        ...

.typings/mflux/models/common/lora/mapping/lora_loader.pyi

@@ -1,27 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from collections.abc import Callable
-from dataclasses import dataclass
-
-import mlx.core as mx
-import mlx.nn as nn
-from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
-
-@dataclass
-class PatternMatch:
-    source_pattern: str
-    target_path: str
-    matrix_name: str
-    transpose: bool
-    transform: Callable[[mx.array], mx.array] | None = ...
-
-class LoRALoader:
-    @staticmethod
-    def load_and_apply_lora(
-        lora_mapping: list[LoRATarget],
-        transformer: nn.Module,
-        lora_paths: list[str] | None = ...,
-        lora_scales: list[float] | None = ...,
-    ) -> tuple[list[str], list[float]]: ...

.typings/mflux/models/common/lora/mapping/lora_mapping.pyi

@@ -1,22 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import List, Protocol
-
-import mlx.core as mx
-
-@dataclass
-class LoRATarget:
-    model_path: str
-    possible_up_patterns: List[str]
-    possible_down_patterns: List[str]
-    possible_alpha_patterns: List[str] = ...
-    up_transform: Callable[[mx.array], mx.array] | None = ...
-    down_transform: Callable[[mx.array], mx.array] | None = ...
-
-class LoRAMapping(Protocol):
-    @staticmethod
-    def get_mapping() -> List[LoRATarget]: ...

.typings/mflux/models/common/lora/mapping/lora_saver.pyi

@@ -1,9 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.nn as nn
-
-class LoRASaver:
-    @staticmethod
-    def bake_and_strip_lora(module: nn.Module) -> nn.Module: ...

.typings/mflux/models/common/lora/mapping/lora_transforms.pyi

@@ -1,35 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-import mlx.core as mx
-
-class LoraTransforms:
-    @staticmethod
-    def split_q_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_k_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_v_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_q_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_k_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_v_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_q_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_k_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_v_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_mlp_up(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_q_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_k_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_v_down(tensor: mx.array) -> mx.array: ...
-    @staticmethod
-    def split_single_mlp_down(tensor: mx.array) -> mx.array: ...

.typings/mflux/models/common/resolution/__init__.pyi

@@ -1,17 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from mflux.models.common.resolution.config_resolution import ConfigResolution
-from mflux.models.common.resolution.lora_resolution import LoraResolution
-from mflux.models.common.resolution.path_resolution import PathResolution
-from mflux.models.common.resolution.quantization_resolution import (
-    QuantizationResolution,
-)
-
-__all__ = [
-    "ConfigResolution",
-    "LoraResolution",
-    "PathResolution",
-    "QuantizationResolution",
-]

.typings/mflux/models/common/resolution/actions.pyi

@@ -1,38 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from enum import Enum
-from typing import NamedTuple
-
-class QuantizationAction(Enum):
-    NONE = ...
-    STORED = ...
-    REQUESTED = ...
-
-class PathAction(Enum):
-    LOCAL = ...
-    HUGGINGFACE_CACHED = ...
-    HUGGINGFACE = ...
-    ERROR = ...
-
-class LoraAction(Enum):
-    LOCAL = ...
-    REGISTRY = ...
-    HUGGINGFACE_COLLECTION_CACHED = ...
-    HUGGINGFACE_COLLECTION = ...
-    HUGGINGFACE_REPO_CACHED = ...
-    HUGGINGFACE_REPO = ...
-    ERROR = ...
-
-class ConfigAction(Enum):
-    EXACT_MATCH = ...
-    EXPLICIT_BASE = ...
-    INFER_SUBSTRING = ...
-    ERROR = ...
-
-class Rule(NamedTuple):
-    priority: int
-    name: str
-    check: str
-    action: QuantizationAction | PathAction | LoraAction | ConfigAction

.typings/mflux/models/common/resolution/config_resolution.pyi

@@ -1,15 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from typing import TYPE_CHECKING
-
-from mflux.models.common.config.model_config import ModelConfig
-
-if TYPE_CHECKING: ...
-logger = ...
-
-class ConfigResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(model_name: str, base_model: str | None = ...) -> ModelConfig: ...

.typings/mflux/models/common/resolution/lora_resolution.pyi

@@ -1,21 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-
-logger = ...
-
-class LoraResolution:
-    RULES = ...
-    _registry: dict[str, Path] = ...
-    @staticmethod
-    def resolve(path: str) -> str: ...
-    @staticmethod
-    def resolve_paths(paths: list[str] | None) -> list[str]: ...
-    @staticmethod
-    def resolve_scales(scales: list[float] | None, num_paths: int) -> list[float]: ...
-    @staticmethod
-    def get_registry() -> dict[str, Path]: ...
-    @staticmethod
-    def discover_files(library_paths: list[Path]) -> dict[str, Path]: ...

.typings/mflux/models/common/resolution/path_resolution.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from pathlib import Path
-
-logger = ...
-
-class PathResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(path: str | None, patterns: list[str] | None = ...) -> Path | None: ...

.typings/mflux/models/common/resolution/quantization_resolution.pyi

@@ -1,12 +0,0 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-logger = ...
-
-class QuantizationResolution:
-    RULES = ...
-    @staticmethod
-    def resolve(
-        stored: int | None, requested: int | None
-    ) -> tuple[int | None, str | None]: ...