← back to Exo
Add mlx nn stubs
0bb621b65373240b036e26a639ef45dabe5fd143 · 2025-11-06 03:59:37 -0800 · rltakashige
Files touched
M src/exo/engines/mlx/__init__.pyM src/exo/engines/mlx/auto_parallel.pyM src/exo/engines/mlx/utils_mlx.pyM src/exo/master/tests/test_master.pyM src/exo/shared/global_conn.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/models/model_meta.pyM src/exo/shared/tests/test_node_id_persistence.pyM src/exo/shared/types/api.pyM src/exo/shared/types/worker/commands_runner.pyM src/exo/shared/types/worker/resource_monitor.pyM src/exo/utils/phantom.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/download/huggingface_utils.pyM src/exo/worker/download/impl_shard_downloader.pyM src/exo/worker/main.pyM src/exo/worker/runner/generate.pyM src/exo/worker/runner/runner_supervisor.pyM src/exo/worker/tests/conftest.pyM src/exo/worker/tests/test_plan/test_worker_plan_utils.pyM src/exo/worker/utils/macmon/macmon.pyM typings/mlx/core/__init__.pyiA typings/mlx/nn/__init__.pyiA typings/mlx/nn/init.pyiA typings/mlx/nn/layers/__init__.pyiA typings/mlx/nn/layers/activations.pyiA typings/mlx/nn/layers/base.pyiA typings/mlx/nn/layers/containers.pyiA typings/mlx/nn/layers/convolution.pyiA typings/mlx/nn/layers/convolution_transpose.pyiA typings/mlx/nn/layers/distributed.pyiA typings/mlx/nn/layers/dropout.pyiA typings/mlx/nn/layers/embedding.pyiA typings/mlx/nn/layers/linear.pyiA typings/mlx/nn/layers/normalization.pyiA typings/mlx/nn/layers/pooling.pyiA typings/mlx/nn/layers/positional_encoding.pyiA typings/mlx/nn/layers/quantized.pyiA typings/mlx/nn/layers/recurrent.pyiA typings/mlx/nn/layers/transformer.pyiA typings/mlx/nn/layers/upsample.pyiA typings/mlx/nn/losses.pyiA typings/mlx/nn/utils.pyiA typings/mlx/utils.pyi
Diff
commit 0bb621b65373240b036e26a639ef45dabe5fd143
Author: rltakashige <rl.takashige@gmail.com>
Date: Thu Nov 6 03:59:37 2025 -0800
Add mlx nn stubs
---
src/exo/engines/mlx/__init__.py | 20 +-
src/exo/engines/mlx/auto_parallel.py | 39 +-
src/exo/engines/mlx/utils_mlx.py | 33 +-
src/exo/master/tests/test_master.py | 4 +-
src/exo/shared/global_conn.py | 3 +-
src/exo/shared/models/model_cards.py | 4 +-
src/exo/shared/models/model_meta.py | 18 +-
src/exo/shared/tests/test_node_id_persistence.py | 3 +-
src/exo/shared/types/api.py | 6 +-
src/exo/shared/types/worker/commands_runner.py | 2 +-
src/exo/shared/types/worker/resource_monitor.py | 6 +-
src/exo/utils/phantom.py | 5 +-
src/exo/worker/download/download_utils.py | 34 +-
src/exo/worker/download/huggingface_utils.py | 14 +-
src/exo/worker/download/impl_shard_downloader.py | 16 +-
src/exo/worker/main.py | 4 +-
src/exo/worker/runner/generate.py | 53 +--
src/exo/worker/runner/runner_supervisor.py | 4 +-
src/exo/worker/tests/conftest.py | 16 +-
.../tests/test_plan/test_worker_plan_utils.py | 8 +-
src/exo/worker/utils/macmon/macmon.py | 39 +-
typings/mlx/core/__init__.pyi | 6 +-
typings/mlx/nn/__init__.pyi | 9 +
typings/mlx/nn/init.pyi | 284 +++++++++++
typings/mlx/nn/layers/__init__.pyi | 20 +
typings/mlx/nn/layers/activations.pyi | 523 +++++++++++++++++++++
typings/mlx/nn/layers/base.pyi | 393 ++++++++++++++++
typings/mlx/nn/layers/containers.pyi | 21 +
typings/mlx/nn/layers/convolution.pyi | 116 +++++
typings/mlx/nn/layers/convolution_transpose.pyi | 119 +++++
typings/mlx/nn/layers/distributed.pyi | 227 +++++++++
typings/mlx/nn/layers/dropout.pyi | 65 +++
typings/mlx/nn/layers/embedding.pyi | 34 ++
typings/mlx/nn/layers/linear.pyi | 76 +++
typings/mlx/nn/layers/normalization.pyi | 194 ++++++++
typings/mlx/nn/layers/pooling.pyi | 242 ++++++++++
typings/mlx/nn/layers/positional_encoding.pyi | 80 ++++
typings/mlx/nn/layers/quantized.pyi | 125 +++++
typings/mlx/nn/layers/recurrent.pyi | 113 +++++
typings/mlx/nn/layers/transformer.pyi | 168 +++++++
typings/mlx/nn/layers/upsample.pyi | 87 ++++
typings/mlx/nn/losses.pyi | 419 +++++++++++++++++
typings/mlx/nn/utils.pyi | 73 +++
typings/mlx/utils.pyi | 182 +++++++
44 files changed, 3734 insertions(+), 173 deletions(-)
diff --git a/src/exo/engines/mlx/__init__.py b/src/exo/engines/mlx/__init__.py
index 716ee0b9..8c0c8fa3 100644
--- a/src/exo/engines/mlx/__init__.py
+++ b/src/exo/engines/mlx/__init__.py
@@ -1,9 +1,9 @@
-from typing import Optional
+from typing import Any
from mlx_lm.models.cache import KVCache
import mlx.core as mx
-import mlx.nn as nn # type: ignore
+import mlx.nn as nn
# These are wrapper functions to fix the fact that mlx is not strongly typed in the same way that EXO is.
# For example - MLX has no guarantee of the interface that nn.Module will expose. But we need a guarantee that it has a __call__() function
@@ -12,7 +12,12 @@ import mlx.nn as nn # type: ignore
class Model(nn.Module):
layers: list[nn.Module]
- def __call__(self, x: mx.array, cache: Optional[list[KVCache]]) -> mx.array: ...
+ def __call__(
+ self,
+ x: mx.array,
+ cache: list[KVCache] | None,
+ input_embeddings: mx.array | None = None,
+ ) -> mx.array: ...
class Detokenizer:
@@ -25,8 +30,15 @@ class Detokenizer:
class TokenizerWrapper:
- bos_token: Optional[str]
+ bos_token: str | None
eos_token_ids: list[int]
detokenizer: Detokenizer
def encode(self, text: str, add_special_tokens: bool = True) -> list[int]: ...
+
+ def apply_chat_template(
+ self,
+ messages_dicts: list[dict[str, Any]],
+ tokenize: bool = False,
+ add_generation_prompt: bool = True,
+ ) -> str: ...
diff --git a/src/exo/engines/mlx/auto_parallel.py b/src/exo/engines/mlx/auto_parallel.py
index 7db609d3..78109325 100644
--- a/src/exo/engines/mlx/auto_parallel.py
+++ b/src/exo/engines/mlx/auto_parallel.py
@@ -1,6 +1,6 @@
from abc import ABC, abstractmethod
from functools import partial
-from typing import TYPE_CHECKING, Protocol, cast, override
+from typing import TYPE_CHECKING, Callable, Protocol, cast, override
from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
@@ -9,16 +9,16 @@ from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
import mlx.core as mx
-import mlx.nn as nn # pyright: ignore[reportMissingTypeStubs]
+import mlx.nn as nn
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
TensorShardMetadata,
)
-from mlx.nn.layers.distributed import ( # type: ignore
- shard_inplace, # type: ignore
- shard_linear, # type: ignore
- sum_gradients, # type: ignore
+from mlx.nn.layers.distributed import (
+ shard_inplace,
+ shard_linear,
+ sum_gradients,
)
@@ -226,18 +226,18 @@ class TensorParallelisationStrategy(ParallelisationShardStrategy):
class TensorParallelShardingStrategy(ABC):
def __init__(
self,
- group, # type: ignore
- all_to_sharded_linear, # type: ignore
- sharded_to_all_linear, # type: ignore
- all_to_sharded_linear_in_place, # type: ignore
- sharded_to_all_linear_in_place, # type: ignore
+ group: mx.distributed.Group,
+ all_to_sharded_linear: Callable[..., nn.Linear],
+ sharded_to_all_linear: Callable[..., nn.Linear],
+ all_to_sharded_linear_in_place: Callable[..., None],
+ sharded_to_all_linear_in_place: Callable[..., None],
):
self.all_to_sharded_linear = all_to_sharded_linear
self.sharded_to_all_linear = sharded_to_all_linear
self.all_to_sharded_linear_in_place = all_to_sharded_linear_in_place
self.sharded_to_all_linear_in_place = sharded_to_all_linear_in_place
- self.group = group or mx.distributed.init() # type: ignore
- self.N = cast(int, group.size()) # type: ignore
+ self.group = group or mx.distributed.init()
+ self.N = group.size()
@abstractmethod
def shard_model(self, model: nn.Module) -> nn.Module: ...
@@ -268,6 +268,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
for layer in model.layers:
# Shard the self attention
if layer.self_attn.q_lora_rank is None: # pyright: ignore[reportUnnecessaryComparison]
+ # Unfortunately, q_lora_rank can be None despite typing hints.
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
)
@@ -297,7 +298,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
layer.mlp = ShardedDeepseekV3MoE(layer.mlp) # type: ignore
- layer.mlp.sharding_group = self.group # type: ignore
+ layer.mlp.sharding_group = self.group
return model
@@ -309,8 +310,8 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
def __call__(self, x: mx.array) -> mx.array:
if self.sharding_group is not None:
- x = sum_gradients(self.sharding_group)(x) # type: ignore
- y = self.original_layer.__call__(x) # type: ignore
+ x = sum_gradients(self.sharding_group)(x)
+ y = self.original_layer.__call__(x)
if self.sharding_group is not None:
y = mx.distributed.all_sum(y, group=self.sharding_group)
return y
@@ -335,7 +336,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
layer.mlp = ShardedQwenMoE(layer.mlp) # type: ignore
- layer.mlp.sharding_group = self.group # type:ignore
+ layer.mlp.sharding_group = self.group
# Shard the MLP
else:
@@ -353,8 +354,8 @@ class ShardedQwenMoE(CustomMlxLayer):
def __call__(self, x: mx.array) -> mx.array:
if self.sharding_group is not None:
- x = sum_gradients(self.sharding_group)(x) # type: ignore
- y = self.original_layer.__call__(x) # type: ignore
+ x = sum_gradients(self.sharding_group)(x)
+ y = self.original_layer.__call__(x)
if self.sharding_group is not None:
y = mx.distributed.all_sum(y, group=self.sharding_group)
return y
diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py
index eb82246c..d1216e73 100644
--- a/src/exo/engines/mlx/utils_mlx.py
+++ b/src/exo/engines/mlx/utils_mlx.py
@@ -3,11 +3,10 @@ import concurrent.futures
import os
import resource
from asyncio import AbstractEventLoop
-from typing import Any, Callable, Optional
+from typing import Any, Callable, cast
from mlx_lm.models.cache import KVCache
from mlx_lm.sample_utils import make_sampler
-from mlx_lm.tokenizer_utils import TokenizerWrapper as _TokenizerWrapper
try:
from mlx_lm.tokenizer_utils import load_tokenizer # type: ignore
@@ -17,7 +16,7 @@ from mlx_lm.utils import load_model # type: ignore
from pydantic import RootModel
import mlx.core as mx
-import mlx.nn as nn # pyright: ignore[reportMissingTypeStubs]
+import mlx.nn as nn
from exo.engines.mlx import Model, TokenizerWrapper
from exo.engines.mlx.auto_parallel import (
IdentityLayer,
@@ -150,15 +149,12 @@ def initialize_mlx(
mlx_ibv_coordinator=mlx_ibv_coordinator,
)
- # set_wired_limit_for_model(get_weights_size(model_shard_meta))
-
- # Determine world size from either hosts or mlx_ibv_devices
-
sampler: Callable[[mx.array], mx.array] = make_sampler(temp=0.7)
model, tokenizer = shard_and_load(model_shard_meta, group=group)
+ model = cast(Model, model)
- return model, tokenizer, sampler, group # type: ignore[return-value]
+ return model, tokenizer, sampler, group
def shard_and_load(
@@ -176,12 +172,9 @@ def shard_and_load(
assert isinstance(model, nn.Module)
tokenizer = load_tokenizer(model_path) # type: ignore
- assert isinstance(tokenizer, _TokenizerWrapper)
+ tokenizer = cast(TokenizerWrapper, tokenizer)
- if group:
- runner_print(f"Group size: {group.size()}, group rank: {group.rank()}")
- else:
- runner_print("!!! No group")
+ runner_print(f"Group size: {group.size()}, group rank: {group.rank()}")
match model_shard_meta.strategy:
case "auto":
@@ -199,13 +192,13 @@ def shard_and_load(
runner_print(f"Model after auto_parallel: {str(model)}")
- mx.eval(model.parameters()) # type: ignore
+ mx.eval(model.parameters())
mx.eval(model)
# Synchronize processes before generation to avoid timeout
mx_barrier(group)
- return model, tokenizer # type: ignore
+ return model, tokenizer
async def apply_chat_template(
@@ -217,10 +210,10 @@ async def apply_chat_template(
# Now we can properly access the messages
messages = chat_task_data.messages
- messages_dicts = [msg.model_dump() for msg in messages]
+ messages_dicts: list[dict[str, Any]] = [msg.model_dump() for msg in messages]
# Filter out None values, keeping relevant keys for the model
- formatted_messages = []
+ formatted_messages: list[dict[str, Any]] = []
for message in messages_dicts:
filtered_message: dict[str, Any] = {
k: v
@@ -235,13 +228,13 @@ async def apply_chat_template(
# If neither content nor thinking is present, skip this message
continue
- formatted_messages.append(filtered_message) # type: ignore
+ formatted_messages.append(filtered_message)
messages_dicts = formatted_messages
prompt: str = await loop.run_in_executor(
executor=mlx_executor,
- func=lambda: tokenizer.apply_chat_template( # type: ignore
+ func=lambda: tokenizer.apply_chat_template(
messages_dicts,
tokenize=False,
add_generation_prompt=True,
@@ -276,7 +269,7 @@ class NullKVCache(KVCache):
async def make_kv_cache(
model: Model,
- max_kv_size: Optional[int] = None,
+ max_kv_size: int | None = None,
) -> list[KVCache]:
assert hasattr(model, "layers")
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 62f51e6d..53b3fced 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -1,4 +1,4 @@
-from typing import List, Sequence
+from typing import Sequence
import anyio
import pytest
@@ -44,7 +44,7 @@ async def test_master():
command_sender, co_receiver = channel[ForwarderCommand]()
local_event_sender, le_receiver = channel[ForwarderEvent]()
- all_events: List[IndexedEvent] = []
+ all_events: list[IndexedEvent] = []
def _get_events() -> Sequence[IndexedEvent]:
orig_events = global_event_receiver.collect()
diff --git a/src/exo/shared/global_conn.py b/src/exo/shared/global_conn.py
index 7ecf5928..01bfc203 100644
--- a/src/exo/shared/global_conn.py
+++ b/src/exo/shared/global_conn.py
@@ -3,7 +3,6 @@
import asyncio
import threading
from multiprocessing.connection import Connection
-from typing import Optional
from exo.shared.types.worker.commands_runner import (
RunnerMessage,
@@ -54,7 +53,7 @@ class AsyncConnection[SendT, RecvT]:
self._conn.close()
-_conn: Optional[AsyncConnection[RunnerResponse, RunnerMessage]] = None
+_conn: AsyncConnection[RunnerResponse, RunnerMessage] | None = None
def set_conn(c: AsyncConnection[RunnerResponse, RunnerMessage]) -> None:
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index c58d3fa1..8fc85c48 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -1,5 +1,3 @@
-from typing import List
-
from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.utils.pydantic_ext import CamelCaseModel
@@ -10,7 +8,7 @@ class ModelCard(CamelCaseModel):
model_id: str
name: str
description: str
- tags: List[str]
+ tags: list[str]
metadata: ModelMetadata
diff --git a/src/exo/shared/models/model_meta.py b/src/exo/shared/models/model_meta.py
index b9bb470a..24da284c 100644
--- a/src/exo/shared/models/model_meta.py
+++ b/src/exo/shared/models/model_meta.py
@@ -1,4 +1,4 @@
-from typing import Annotated, Dict, Optional
+from typing import Annotated
import aiofiles
import aiofiles.os as aios
@@ -19,14 +19,12 @@ class ConfigData(BaseModel):
model_config = {"extra": "ignore"} # Allow unknown fields
# Common field names for number of layers across different architectures
- num_hidden_layers: Optional[Annotated[int, Field(ge=0)]] = None
- num_layers: Optional[Annotated[int, Field(ge=0)]] = None
- n_layer: Optional[Annotated[int, Field(ge=0)]] = None
- n_layers: Optional[Annotated[int, Field(ge=0)]] = None # Sometimes used
- num_decoder_layers: Optional[Annotated[int, Field(ge=0)]] = (
- None # Transformer models
- )
- decoder_layers: Optional[Annotated[int, Field(ge=0)]] = None # Some architectures
+ num_hidden_layers: Annotated[int, Field(ge=0)] | None = None
+ num_layers: Annotated[int, Field(ge=0)] | None = None
+ n_layer: Annotated[int, Field(ge=0)] | None = None
+ n_layers: Annotated[int, Field(ge=0)] | None = None # Sometimes used
+ num_decoder_layers: Annotated[int, Field(ge=0)] | None = None # Transformer models
+ decoder_layers: Annotated[int, Field(ge=0)] | None = None # Some architectures
@property
def layer_count(self) -> int:
@@ -92,7 +90,7 @@ async def get_safetensors_size(model_id: str) -> Memory:
return Memory.from_bytes(info.safetensors.total)
-_model_meta_cache: Dict[str, ModelMetadata] = {}
+_model_meta_cache: dict[str, ModelMetadata] = {}
async def get_model_meta(model_id: str) -> ModelMetadata:
diff --git a/src/exo/shared/tests/test_node_id_persistence.py b/src/exo/shared/tests/test_node_id_persistence.py
index 4633ab90..cdcf19ca 100644
--- a/src/exo/shared/tests/test_node_id_persistence.py
+++ b/src/exo/shared/tests/test_node_id_persistence.py
@@ -7,7 +7,6 @@ from multiprocessing.process import BaseProcess
from multiprocessing.queues import Queue as QueueT
from multiprocessing.synchronize import Event as EventT
from multiprocessing.synchronize import Semaphore as SemaphoreT
-from typing import Optional
from pytest import LogCaptureFixture
@@ -64,7 +63,7 @@ def _get_keypair_concurrent(num_procs: int) -> bytes:
# check that the input/output order match, and that
# all subprocesses end up reading the same file
logging.info(msg="PARENT: Checking consistency")
- keypair: Optional[bytes] = None
+ keypair: bytes | None = None
qsize = 0 # cannot use Queue.qsize due to MacOS incompatibility :(
while not queue.empty():
qsize += 1
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 88044379..e5437c4e 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -1,5 +1,5 @@
import time
-from typing import Any, List, Literal
+from typing import Any, Literal
from pydantic import BaseModel, Field
@@ -20,12 +20,12 @@ class ModelListModel(BaseModel):
name: str = Field(default="")
description: str = Field(default="")
context_length: int = Field(default=0)
- tags: List[str] = Field(default=[])
+ tags: list[str] = Field(default=[])
class ModelList(BaseModel):
object: str = "list"
- data: List[ModelListModel]
+ data: list[ModelListModel]
class ChatCompletionMessage(BaseModel):
diff --git a/src/exo/shared/types/worker/commands_runner.py b/src/exo/shared/types/worker/commands_runner.py
index 465d3887..0a873f8a 100644
--- a/src/exo/shared/types/worker/commands_runner.py
+++ b/src/exo/shared/types/worker/commands_runner.py
@@ -43,7 +43,7 @@ class TokenizedResponse(BaseRunnerResponse):
class GenerationResponse(BaseRunnerResponse):
text: str
token: int
- # logprobs: Optional[list[float]] = None # too big. we can change to be top-k
+ # logprobs: list[float] | None = None # too big. we can change to be top-k
finish_reason: FinishReason | None = None
diff --git a/src/exo/shared/types/worker/resource_monitor.py b/src/exo/shared/types/worker/resource_monitor.py
index 8a1f3349..b351963c 100644
--- a/src/exo/shared/types/worker/resource_monitor.py
+++ b/src/exo/shared/types/worker/resource_monitor.py
@@ -1,7 +1,7 @@
import asyncio
from abc import ABC, abstractmethod
from collections.abc import Coroutine
-from typing import Callable, List, Set
+from typing import Callable
from exo.shared.types.profiling import (
MemoryPerformanceProfile,
@@ -23,8 +23,8 @@ class MemoryResourceCollector(ResourceCollector):
class ResourceMonitor:
- data_collectors: List[ResourceCollector]
- effect_handlers: Set[
+ data_collectors: list[ResourceCollector]
+ effect_handlers: set[
Callable[[SystemPerformanceProfile | MemoryPerformanceProfile], None]
]
diff --git a/src/exo/utils/phantom.py b/src/exo/utils/phantom.py
index 4fe62afb..72e33442 100644
--- a/src/exo/utils/phantom.py
+++ b/src/exo/utils/phantom.py
@@ -1,13 +1,10 @@
-from typing import Optional
-
-
class _PhantomData[*T]:
"""
Internal machinery of the phantom data - it stores nothing.
"""
-type PhantomData[*T] = Optional[_PhantomData[*T]]
+type PhantomData[*T] = _PhantomData[*T] | None
"""
Allows you to use generics in functions without storing anything of that generic type.
Just use `None` and you'll be fine
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index 2a7f6cf1..c179b921 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -6,7 +6,7 @@ import time
import traceback
from datetime import timedelta
from pathlib import Path
-from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Callable, Literal
from urllib.parse import urljoin
import aiofiles
@@ -39,8 +39,8 @@ class ModelSafetensorsIndexMetadata(BaseModel):
class ModelSafetensorsIndex(BaseModel):
- metadata: Optional[ModelSafetensorsIndexMetadata]
- weight_map: Dict[str, str]
+ metadata: ModelSafetensorsIndexMetadata | None
+ weight_map: dict[str, str]
class FileListEntry(BaseModel):
@@ -76,7 +76,7 @@ class RepoDownloadProgress(BaseModel):
overall_speed: float
overall_eta: timedelta
status: Literal["not_started", "in_progress", "complete"]
- file_progress: Dict[str, RepoFileDownloadProgress] = Field(default_factory=dict)
+ file_progress: dict[str, RepoFileDownloadProgress] = Field(default_factory=dict)
model_config = ConfigDict(frozen=True)
@@ -162,7 +162,7 @@ async def delete_model(repo_id: str) -> bool:
return True
-async def seed_models(seed_dir: Union[str, Path]):
+async def seed_models(seed_dir: str | Path):
"""Move model in resources folder of app to .cache/huggingface/hub"""
source_dir = Path(seed_dir)
dest_dir = await ensure_models_dir()
@@ -181,7 +181,7 @@ async def seed_models(seed_dir: Union[str, Path]):
async def fetch_file_list_with_cache(
repo_id: str, revision: str = "main", recursive: bool = False
-) -> List[FileListEntry]:
+) -> list[FileListEntry]:
target_dir = (
(await ensure_models_dir()) / "caches" / str(repo_id).replace("/", "--")
)
@@ -191,17 +191,17 @@ async def fetch_file_list_with_cache(
)
if await aios.path.exists(cache_file):
async with aiofiles.open(cache_file, "r") as f:
- return TypeAdapter(List[FileListEntry]).validate_json(await f.read())
+ return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
file_list = await fetch_file_list_with_retry(repo_id, revision, recursive=recursive)
await aios.makedirs(cache_file.parent, exist_ok=True)
async with aiofiles.open(cache_file, "w") as f:
- await f.write(TypeAdapter(List[FileListEntry]).dump_json(file_list).decode())
+ await f.write(TypeAdapter(list[FileListEntry]).dump_json(file_list).decode())
return file_list
async def fetch_file_list_with_retry(
repo_id: str, revision: str = "main", path: str = "", recursive: bool = False
-) -> List[FileListEntry]:
+) -> list[FileListEntry]:
n_attempts = 30
for attempt in range(n_attempts):
try:
@@ -217,7 +217,7 @@ async def fetch_file_list_with_retry(
async def _fetch_file_list(
repo_id: str, revision: str = "main", path: str = "", recursive: bool = False
-) -> List[FileListEntry]:
+) -> list[FileListEntry]:
api_url = f"{get_hf_endpoint()}/api/models/{repo_id}/tree/{revision}"
url = f"{api_url}/{path}" if path else api_url
@@ -286,7 +286,7 @@ async def calc_hash(path: Path, hash_type: Literal["sha1", "sha256"] = "sha1") -
async def file_meta(
repo_id: str, revision: str, path: str, redirected_location: str | None = None
-) -> Tuple[int, str]:
+) -> tuple[int, str]:
url = (
urljoin(f"{get_hf_endpoint()}/{repo_id}/resolve/{revision}/", path)
if redirected_location is None
@@ -405,7 +405,7 @@ def calculate_repo_progress(
shard: ShardMetadata,
repo_id: str,
revision: str,
- file_progress: Dict[str, RepoFileDownloadProgress],
+ file_progress: dict[str, RepoFileDownloadProgress],
all_start_time: float,
) -> RepoDownloadProgress:
all_total_bytes = sum((p.total.in_bytes for p in file_progress.values()), 0)
@@ -451,7 +451,7 @@ def calculate_repo_progress(
)
-async def get_weight_map(repo_id: str, revision: str = "main") -> Dict[str, str]:
+async def get_weight_map(repo_id: str, revision: str = "main") -> dict[str, str]:
target_dir = (await ensure_models_dir()) / str(repo_id).replace("/", "--")
await aios.makedirs(target_dir, exist_ok=True)
index_file = await download_file_with_retry(
@@ -462,7 +462,7 @@ async def get_weight_map(repo_id: str, revision: str = "main") -> Dict[str, str]
return index_data.weight_map
-async def resolve_allow_patterns(shard: ShardMetadata) -> List[str]:
+async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
try:
weight_map = await get_weight_map(str(shard.model_meta.model_id))
return get_allow_patterns(weight_map, shard)
@@ -484,7 +484,7 @@ async def get_downloaded_size(path: Path) -> int:
async def download_progress_for_local_path(
repo_id: str, shard: ShardMetadata, local_path: Path
) -> RepoDownloadProgress:
- file_progress: Dict[str, RepoFileDownloadProgress] = {}
+ file_progress: dict[str, RepoFileDownloadProgress] = {}
total_files = 0
total_bytes = 0
@@ -533,7 +533,7 @@ async def download_shard(
on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
max_parallel_downloads: int = 8,
skip_download: bool = False,
- allow_patterns: List[str] | None = None,
+ allow_patterns: list[str] | None = None,
) -> tuple[Path, RepoDownloadProgress]:
if not skip_download:
logger.info(f"Downloading {shard.model_meta.model_id=}")
@@ -568,7 +568,7 @@ async def download_shard(
file_list, allow_patterns=allow_patterns, key=lambda x: x.path
)
)
- file_progress: Dict[str, RepoFileDownloadProgress] = {}
+ file_progress: dict[str, RepoFileDownloadProgress] = {}
def on_progress_wrapper(
file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
diff --git a/src/exo/worker/download/huggingface_utils.py b/src/exo/worker/download/huggingface_utils.py
index 2e3df1b8..fbf711e1 100644
--- a/src/exo/worker/download/huggingface_utils.py
+++ b/src/exo/worker/download/huggingface_utils.py
@@ -1,7 +1,7 @@
import os
from fnmatch import fnmatch
from pathlib import Path
-from typing import Callable, Dict, Generator, Iterable, List, Optional, TypeVar, Union
+from typing import Callable, Generator, Iterable, TypeVar
import aiofiles
import aiofiles.os as aios
@@ -15,9 +15,9 @@ T = TypeVar("T")
def filter_repo_objects(
items: Iterable[T],
*,
- allow_patterns: Optional[Union[List[str], str]] = None,
- ignore_patterns: Optional[Union[List[str], str]] = None,
- key: Optional[Callable[[T], str]] = None,
+ allow_patterns: list[str] | str | None = None,
+ ignore_patterns: list[str] | str | None = None,
+ key: Callable[[T], str] | None = None,
) -> Generator[T, None, None]:
if isinstance(allow_patterns, str):
allow_patterns = [allow_patterns]
@@ -69,7 +69,7 @@ def get_hf_home() -> Path:
return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
-async def get_hf_token() -> Optional[str]:
+async def get_hf_token() -> str | None:
"""Retrieve the Hugging Face token from the user's HF_HOME directory."""
token_path = get_hf_home() / "token"
if await aios.path.exists(token_path):
@@ -86,7 +86,7 @@ async def get_auth_headers() -> dict[str, str]:
return {}
-def extract_layer_num(tensor_name: str) -> Optional[int]:
+def extract_layer_num(tensor_name: str) -> int | None:
# This is a simple example and might need to be adjusted based on the actual naming convention
parts = tensor_name.split(".")
for part in parts:
@@ -95,7 +95,7 @@ def extract_layer_num(tensor_name: str) -> Optional[int]:
return None
-def get_allow_patterns(weight_map: Dict[str, str], shard: ShardMetadata) -> List[str]:
+def get_allow_patterns(weight_map: dict[str, str], shard: ShardMetadata) -> list[str]:
default_patterns = set(["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"])
shard_specific_patterns: set[str] = set()
if weight_map:
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/worker/download/impl_shard_downloader.py
index 67f3236c..a00ac5a7 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -1,6 +1,6 @@
import asyncio
from pathlib import Path
-from typing import AsyncIterator, Callable, Dict, List, Optional
+from typing import AsyncIterator, Callable
from exo.shared.models.model_cards import MODEL_CARDS
from exo.shared.models.model_meta import get_model_meta
@@ -18,7 +18,7 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
)
-async def build_base_shard(model_id: str) -> Optional[ShardMetadata]:
+async def build_base_shard(model_id: str) -> ShardMetadata:
model_meta = await get_model_meta(model_id)
# print(f"build_base_shard {model_id=} {model_meta=}")
return PipelineShardMetadata(
@@ -31,10 +31,8 @@ async def build_base_shard(model_id: str) -> Optional[ShardMetadata]:
)
-async def build_full_shard(model_id: str) -> Optional[PipelineShardMetadata]:
+async def build_full_shard(model_id: str) -> PipelineShardMetadata | None:
base_shard = await build_base_shard(model_id)
- if base_shard is None:
- return None
return PipelineShardMetadata(
model_meta=base_shard.model_meta,
device_rank=base_shard.device_rank,
@@ -48,7 +46,7 @@ async def build_full_shard(model_id: str) -> Optional[PipelineShardMetadata]:
class SingletonShardDownloader(ShardDownloader):
def __init__(self, shard_downloader: ShardDownloader):
self.shard_downloader = shard_downloader
- self.active_downloads: Dict[ShardMetadata, asyncio.Task[Path]] = {}
+ self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
def on_progress(
self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
@@ -83,7 +81,7 @@ class SingletonShardDownloader(ShardDownloader):
class CachedShardDownloader(ShardDownloader):
def __init__(self, shard_downloader: ShardDownloader):
self.shard_downloader = shard_downloader
- self.cache: Dict[tuple[str, ShardMetadata], Path] = {}
+ self.cache: dict[tuple[str, ShardMetadata], Path] = {}
def on_progress(
self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
@@ -117,7 +115,7 @@ class CachedShardDownloader(ShardDownloader):
class ResumableShardDownloader(ShardDownloader):
def __init__(self, max_parallel_downloads: int = 8):
self.max_parallel_downloads = max_parallel_downloads
- self.on_progress_callbacks: List[
+ self.on_progress_callbacks: list[
Callable[[ShardMetadata, RepoDownloadProgress], None]
] = []
@@ -152,7 +150,7 @@ class ResumableShardDownloader(ShardDownloader):
# print("get_shard_download_status")
async def _status_for_model(
model_id: str,
- ) -> Optional[tuple[Path, RepoDownloadProgress]]:
+ ) -> tuple[Path, RepoDownloadProgress] | None:
"""Helper coroutine that builds the shard for a model and gets its download status."""
shard = await build_full_shard(model_id)
if shard is None:
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 4f48aedb..1091c807 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -3,7 +3,7 @@ import time
from asyncio import Queue
from functools import partial
from random import random
-from typing import AsyncGenerator, Optional
+from typing import AsyncGenerator
import anyio
from anyio import CancelScope, create_task_group
@@ -430,7 +430,7 @@ class Worker:
yield RunnerDeleted(runner_id=op.runner_id)
async def _execute_runner_up_op(
- self, op: RunnerUpOp, initialize_timeout: Optional[float] = None
+ self, op: RunnerUpOp, initialize_timeout: float | None = None
) -> AsyncGenerator[Event, None]:
assigned_runner = self.assigned_runners[op.runner_id]
diff --git a/src/exo/worker/runner/generate.py b/src/exo/worker/runner/generate.py
index 3db14141..d1497263 100644
--- a/src/exo/worker/runner/generate.py
+++ b/src/exo/worker/runner/generate.py
@@ -4,7 +4,7 @@ import functools
import time
from collections.abc import AsyncGenerator
from functools import partial
-from typing import Any, Callable, Generator, Optional, Tuple
+from typing import Any, Callable, Generator
import mlx.core as mx
from mlx.core import array
@@ -54,8 +54,8 @@ def generate_step(
max_tokens: int = 256,
sampler: Callable[[mx.array], mx.array],
logits_processors: list[Callable[[mx.array, mx.array], mx.array]] | None = None,
- max_kv_size: Optional[int] = None,
- prompt_cache: Optional[list[KVCache]] = None,
+ max_kv_size: int | None = None,
+ prompt_cache: list[KVCache] | None = None,
prefill_step_size: int = 2048,
kv_bits: int | None = None,
kv_group_size: int = 64,
@@ -63,7 +63,7 @@ def generate_step(
prompt_progress_callback: Callable[[int, int], None] | None = None,
input_embeddings: mx.array | None = None,
group: mx.distributed.Group | None = None,
-) -> Generator[Tuple[int, mx.array], None, None]:
+) -> Generator[tuple[int, mx.array], None, None]:
"""
A generator producing token ids based on the given prompt from the model.
@@ -74,12 +74,12 @@ def generate_step(
generator. Default: ``256``.
sampler (Callable[mx.array, mx.array]): A sampler for sampling a
token from a vector of log probabilities.
- logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+ logits_processors (list[Callable[[mx.array, mx.array], mx.array]], optional):
A list of functions that take tokens and logits and return the processed
logits. Default: ``None``.
max_kv_size (int, optional): Maximum size of the key-value cache. Old
entries (except the first 4 tokens) will be overwritten.
- prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
+ prompt_cache (list[Any], optional): A pre-computed prompt cache. Note, if
provided, the cache will be updated in place.
prefill_step_size (int): Step size for processing the prompt.
kv_bits (int, optional): Number of bits to use for KV cache quantization.
@@ -93,7 +93,7 @@ def generate_step(
conjunction with prompt tokens. Default: ``None``.
Yields:
- Tuple[int, mx.array]: One token and a vector of log probabilities.
+ tuple[int, mx.array]: One token and a vector of log probabilities.
"""
if input_embeddings is not None:
if len(prompt) > 0 and len(prompt) != len(input_embeddings):
@@ -115,7 +115,7 @@ def generate_step(
max_kv_size=max_kv_size,
)
- prompt_progress_callback = prompt_progress_callback or (lambda *_: None) # type: ignore[type-arg]
+ prompt_progress_callback = prompt_progress_callback or (lambda _, __: None)
quantize_cache_fn = functools.partial(
maybe_quantize_kv_cache,
@@ -128,10 +128,10 @@ def generate_step(
input_tokens: mx.array, input_embeddings: mx.array | None
) -> mx.array:
if input_embeddings is not None:
- return model( # type: ignore[type-arg]
+ return model(
input_tokens,
cache=prompt_cache,
- input_embeddings=input_embeddings, # type: ignore[type-arg]
+ input_embeddings=input_embeddings,
)
else:
return model(input_tokens, cache=prompt_cache)
@@ -209,29 +209,28 @@ def generate_step(
prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
if prompt_processed_tokens > 0:
- runner_print("finished prefil stage.")
+ runner_print("finished prefill stage.")
y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
+ prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
+
mx.async_eval(y, logprobs)
- next_y: array | None = None
- next_logprobs: array | None = None
+ next_y, next_logprobs = _step(y)
+ mx.async_eval(next_y, next_logprobs)
n = 0
+
while True:
- if n != max_tokens:
- assert y is not None
- next_y, next_logprobs = _step(y)
- mx.async_eval(next_y, next_logprobs)
- if n == 0:
- mx.eval(y) # type: ignore[type-arg]
- prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
- if n == max_tokens:
- break
- yield int(y.item()), logprobs # type: ignore
+ mx.eval(y, logprobs)
+ yield int(y.item()), logprobs
+ n += 1
if n % 256 == 0:
mx.clear_cache()
- y, logprobs = next_y, next_logprobs
- n += 1
+ if n == max_tokens:
+ break
+ y, logprobs = next_y, logprobs
+ next_y, next_logprobs = _step(y)
+ mx.async_eval(next_y, next_logprobs)
def stream_generate(
@@ -243,7 +242,7 @@ def stream_generate(
conn: AsyncConnection[RunnerResponse, RunnerMessage] | None,
logits_processors: list[Callable[[mx.array, mx.array], mx.array]] | None = None,
max_kv_size: int | None = None,
- prompt_cache: Optional[list[KVCache]] = None,
+ prompt_cache: list[KVCache] | None = None,
prefill_step_size: int = 2048,
kv_bits: int | None = None,
kv_group_size: int = 64,
@@ -264,7 +263,7 @@ def stream_generate(
detokenizer = tokenizer.detokenizer
- token_generator: Generator[Tuple[int, array], None, None] = generate_step(
+ token_generator: Generator[tuple[int, array], None, None] = generate_step(
prompt_array,
model,
max_tokens=max_tokens,
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index 1923ac96..7a7fe8a9 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -7,7 +7,7 @@ import tempfile
import traceback
from multiprocessing import Process
from multiprocessing.connection import Connection
-from typing import Any, AsyncGenerator, Callable, Coroutine, Optional
+from typing import Any, AsyncGenerator, Callable, Coroutine
import psutil
from loguru import logger
@@ -75,7 +75,7 @@ class RunnerSupervisor:
hosts: list[Host] | None = None,
mlx_ibv_devices: list[list[str | None]] | None = None,
mlx_ibv_coordinator: str | None = None,
- initialize_timeout: Optional[float] = None,
+ initialize_timeout: float | None = None,
) -> "RunnerSupervisor":
"""
Create and initialize a RunnerSupervisor instance.
diff --git a/src/exo/worker/tests/conftest.py b/src/exo/worker/tests/conftest.py
index 3385866f..380a93d9 100644
--- a/src/exo/worker/tests/conftest.py
+++ b/src/exo/worker/tests/conftest.py
@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Callable
import pytest
@@ -109,13 +109,11 @@ def instance(
pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata],
hosts: Callable[[int], list[Host]],
):
- from typing import Optional
-
def _instance(
- instance_id: Optional[InstanceId] = None,
- node_id: Optional[NodeId] = None,
- runner_id: Optional[RunnerId] = None,
- model_id: Optional[ModelId] = None,
+ instance_id: InstanceId | None = None,
+ node_id: NodeId | None = None,
+ runner_id: RunnerId | None = None,
+ model_id: ModelId | None = None,
) -> Instance:
resolved_instance_id = instance_id if instance_id is not None else INSTANCE_1_ID
resolved_node_id = node_id if node_id is not None else NODE_A
@@ -150,8 +148,8 @@ def completion_create_params(user_message: str) -> ChatCompletionTaskParams:
@pytest.fixture
def chat_completion_task(completion_create_params: ChatCompletionTaskParams):
def _chat_completion_task(
- instance_id: Optional[InstanceId] = None,
- task_id: Optional[TaskId] = None,
+ instance_id: InstanceId | None = None,
+ task_id: TaskId | None = None,
user_message: str = "Hello",
) -> ChatCompletionTask:
resolved_instance_id = instance_id if instance_id is not None else INSTANCE_1_ID
diff --git a/src/exo/worker/tests/test_plan/test_worker_plan_utils.py b/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
index bcddf89a..9053df1f 100644
--- a/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
+++ b/src/exo/worker/tests/test_plan/test_worker_plan_utils.py
@@ -1,5 +1,5 @@
from dataclasses import dataclass
-from typing import List, NotRequired, Optional, TypedDict
+from typing import NotRequired, TypedDict
from typing_extensions import Literal
@@ -78,8 +78,8 @@ class PlanTestCase:
description: str
state: State
- in_process_runners: List[InProcessRunner]
- expected_op: Optional[RunnerOp]
+ in_process_runners: list[InProcessRunner]
+ expected_op: RunnerOp | None
def id(self) -> str: # noqa: D401
return self.description.replace(" ", "_")
@@ -229,7 +229,7 @@ def make_test_case(
description: str,
runner_specs: list[RunnerSpecDict],
tasks: list[TaskSpecDict] | None = None,
- expected_op: Optional[RunnerOp] = None,
+ expected_op: RunnerOp | None = None,
instance_id: InstanceId = INSTANCE_1_ID,
instance_status: InstanceStatus = InstanceStatus.Active,
model_id: ModelId = MODEL_A_ID,
diff --git a/src/exo/worker/utils/macmon/macmon.py b/src/exo/worker/utils/macmon/macmon.py
index 1d823c2a..81e949ff 100644
--- a/src/exo/worker/utils/macmon/macmon.py
+++ b/src/exo/worker/utils/macmon/macmon.py
@@ -2,7 +2,6 @@ import asyncio
import platform
import shutil
import subprocess
-from typing import Optional, Tuple
from pydantic import BaseModel, ConfigDict, ValidationError
@@ -43,10 +42,10 @@ def _get_binary_path() -> str:
class MemoryMetrics(BaseModel):
"""Memory-related metrics returned by macmon."""
- ram_total: Optional[int] = None
- ram_usage: Optional[int] = None
- swap_total: Optional[int] = None
- swap_usage: Optional[int] = None
+ ram_total: int | None = None
+ ram_usage: int | None = None
+ swap_total: int | None = None
+ swap_usage: int | None = None
model_config = ConfigDict(extra="ignore")
@@ -54,8 +53,8 @@ class MemoryMetrics(BaseModel):
class TempMetrics(BaseModel):
"""Temperature-related metrics returned by macmon."""
- cpu_temp_avg: Optional[float] = None
- gpu_temp_avg: Optional[float] = None
+ cpu_temp_avg: float | None = None
+ gpu_temp_avg: float | None = None
model_config = ConfigDict(extra="ignore")
@@ -67,19 +66,19 @@ class Metrics(BaseModel):
Unknown fields are ignored for forward-compatibility.
"""
- all_power: Optional[float] = None
- ane_power: Optional[float] = None
- cpu_power: Optional[float] = None
- ecpu_usage: Optional[Tuple[int, float]] = None
- gpu_power: Optional[float] = None
- gpu_ram_power: Optional[float] = None
- gpu_usage: Optional[Tuple[int, float]] = None
- memory: Optional[MemoryMetrics] = None
- pcpu_usage: Optional[Tuple[int, float]] = None
- ram_power: Optional[float] = None
- sys_power: Optional[float] = None
- temp: Optional[TempMetrics] = None
- timestamp: Optional[str] = None
+ all_power: float | None = None
+ ane_power: float | None = None
+ cpu_power: float | None = None
+ ecpu_usage: tuple[int, float] | None = None
+ gpu_power: float | None = None
+ gpu_ram_power: float | None = None
+ gpu_usage: tuple[int, float] | None = None
+ memory: MemoryMetrics | None = None
+ pcpu_usage: tuple[int, float] | None = None
+ ram_power: float | None = None
+ sys_power: float | None = None
+ temp: TempMetrics | None = None
+ timestamp: str | None = None
model_config = ConfigDict(extra="ignore")
diff --git a/typings/mlx/core/__init__.pyi b/typings/mlx/core/__init__.pyi
index e1ffbe29..8edb9832 100644
--- a/typings/mlx/core/__init__.pyi
+++ b/typings/mlx/core/__init__.pyi
@@ -1,10 +1,8 @@
import enum
import pathlib
-import sys
import types
from typing import (
Annotated,
- Any,
Callable,
Literal,
Mapping,
@@ -14,6 +12,7 @@ from typing import (
)
import numpy
+from mlx.nn.layers import Module
from numpy.typing import ArrayLike as _ArrayLike
from . import cuda as cuda
@@ -2609,9 +2608,10 @@ euler_gamma: float = ...
type MX_ARRAY_TREE = (
array
+ | Module
| list[MX_ARRAY_TREE]
| tuple[MX_ARRAY_TREE, ...]
- | Mapping[Any, MX_ARRAY_TREE]
+ | Mapping[str, MX_ARRAY_TREE]
)
def eval(*args: MX_ARRAY_TREE) -> None:
diff --git a/typings/mlx/nn/__init__.pyi b/typings/mlx/nn/__init__.pyi
new file mode 100644
index 00000000..4c999379
--- /dev/null
+++ b/typings/mlx/nn/__init__.pyi
@@ -0,0 +1,9 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from layers import *
+from utils import *
+
+from . import init as init
+from . import losses as losses
diff --git a/typings/mlx/nn/init.pyi b/typings/mlx/nn/init.pyi
new file mode 100644
index 00000000..efa453e6
--- /dev/null
+++ b/typings/mlx/nn/init.pyi
@@ -0,0 +1,284 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Callable, Literal
+
+import mlx.core as mx
+
+def constant(value: float, dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns an array filled with ``value``.
+
+ Args:
+ value (float): The value to fill the array with.
+ dtype (Dtype, optional): The data type of the array. Default:
+ ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns an array with the
+ same shape as the input, filled with ``value``.
+
+ Example:
+
+ >>> init_fn = nn.init.constant(0.5)
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[0.5, 0.5],
+ [0.5, 0.5]], dtype=float32)
+ """
+
+def normal(
+ mean: float = ..., std: float = ..., dtype: mx.Dtype = ...
+) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns samples from a normal distribution.
+
+ Args:
+ mean (float, optional): Mean of the normal distribution. Default:
+ ``0.0``.
+ std (float, optional): Standard deviation of the normal distribution.
+ Default: ``1.0``.
+ dtype (Dtype, optional): The data type of the array. Default:
+ ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns an array with the
+ same shape as the input, filled with samples from a normal distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.normal()
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[-0.982273, -0.534422],
+ [0.380709, 0.0645099]], dtype=float32)
+ """
+
+def uniform(
+ low: float = ..., high: float = ..., dtype: mx.Dtype = ...
+) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns samples from a uniform distribution.
+
+ Args:
+ low (float, optional): The lower bound of the uniform distribution.
+ Default: ``0.0``.
+ high (float, optional): The upper bound of the uniform distribution.
+ Default: ``1.0``
+ dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns an array
+ with the same shape as the input, filled with samples from a uniform
+ distribution
+
+ Example:
+
+ >>> init_fn = nn.init.uniform(low=0, high=1)
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[0.883935, 0.863726],
+ [0.617261, 0.417497]], dtype=float32)
+ """
+
+def identity(dtype: mx.Dtype = ...) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns an identity matrix.
+
+ Args:
+ dtype (Dtype, optional): The data type of the array. Defaults:
+ ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns an identity
+ matrix with the same shape as the input.
+
+ Example:
+
+ >>> init_fn = nn.init.identity()
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[1, 0],
+ [0, 1]], dtype=float32)
+ """
+
+def glorot_normal(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]:
+ r"""A Glorot normal initializer.
+
+ This initializer samples from a normal distribution with a standard
+ deviation computed from the number of input (``fan_in``) and output
+ (``fan_out``) units according to:
+
+ .. math::
+ \sigma = \gamma \sqrt{\frac{2.0}{\text{fan\_in} + \text{fan\_out}}}
+
+ For more details see the original reference: `Understanding the difficulty
+ of training deep feedforward neural networks
+ <https://proceedings.mlr.press/v9/glorot10a.html>`_
+
+ Args:
+ dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+ Returns:
+ Callable[[array, float], array]: An initializer that returns an array
+ with the same shape as the input, filled with samples from the Glorot
+ normal distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.glorot_normal()
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[0.191107, 1.61278],
+ [-0.150594, -0.363207]], dtype=float32)
+ >>> init_fn(mx.zeros((2, 2)), gain=4.0)
+ array([[1.89613, -4.53947],
+ [4.48095, 0.995016]], dtype=float32)
+ """
+
+def glorot_uniform(dtype: mx.Dtype = ...) -> Callable[[mx.array, float], mx.array]:
+ r"""A Glorot uniform initializer.
+
+ This initializer samples from a uniform distribution with a range
+ computed from the number of input (``fan_in``) and output (``fan_out``)
+ units according to:
+
+ .. math::
+ \sigma = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}}
+
+ For more details see the original reference: `Understanding the difficulty
+ of training deep feedforward neural networks
+ <https://proceedings.mlr.press/v9/glorot10a.html>`_
+
+ Args:
+ dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+ Returns:
+ Callable[[array, float], array]: An initializer that returns an array
+ with the same shape as the input, filled with samples from the Glorot
+ uniform distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.glorot_uniform()
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[0.223404, -0.890597],
+ [-0.379159, -0.776856]], dtype=float32)
+ >>> init_fn(mx.zeros((2, 2)), gain=4.0)
+ array([[-1.90041, 3.02264],
+ [-0.912766, 4.12451]], dtype=float32)
+ """
+
+def he_normal(
+ dtype: mx.Dtype = ...,
+) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]:
+ r"""Build a He normal initializer.
+
+ This initializer samples from a normal distribution with a standard
+ deviation computed from the number of input (``fan_in``) or output
+ (``fan_out``) units according to:
+
+ .. math::
+ \sigma = \gamma \frac{1}{\sqrt{\text{fan}}}
+
+ where :math:`\text{fan}` is either the number of input units when the
+ ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
+ ``"fan_out"``.
+
+ For more details see the original reference: `Delving Deep into Rectifiers:
+ Surpassing Human-Level Performance on ImageNet Classification
+ <https://arxiv.org/abs/1502.01852>`_
+
+ Args:
+ dtype (Dtype, optional): The data type of the array. Defaults to mx.float32.
+
+ Returns:
+ Callable[[array, str, float], array]: An initializer that returns an
+ array with the same shape as the input, filled with samples from the He
+ normal distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.he_normal()
+ >>> init_fn(mx.zeros((2, 2))) # uses fan_in
+ array([[-1.25211, 0.458835],
+ [-0.177208, -0.0137595]], dtype=float32)
+ >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
+ array([[5.6967, 4.02765],
+ [-4.15268, -2.75787]], dtype=float32)
+ """
+
+def he_uniform(
+ dtype: mx.Dtype = ...,
+) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]:
+ r"""A He uniform (Kaiming uniform) initializer.
+
+ This initializer samples from a uniform distribution with a range
+ computed from the number of input (``fan_in``) or output (``fan_out``)
+ units according to:
+
+ .. math::
+
+ \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}}
+
+ where :math:`\text{fan}` is either the number of input units when the
+ ``mode`` is ``"fan_in"`` or output units when the ``mode`` is
+ ``"fan_out"``.
+
+ For more details see the original reference: `Delving Deep into Rectifiers:
+ Surpassing Human-Level Performance on ImageNet Classification
+ <https://arxiv.org/abs/1502.01852>`_
+
+
+ Args:
+ dtype (Dtype, optional): The data type of the array. Default: ``float32``.
+
+ Returns:
+ Callable[[array, str, float], array]: An initializer that returns an
+ array with the same shape as the input, filled with samples from the
+ He uniform distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.he_uniform()
+ >>> init_fn(mx.zeros((2, 2))) # uses fan_in
+ array([[0.0300242, -0.0184009],
+ [0.793615, 0.666329]], dtype=float32)
+ >>> init_fn(mx.zeros((2, 2)), mode="fan_out", gain=5)
+ array([[-1.64331, -2.16506],
+ [1.08619, 5.79854]], dtype=float32)
+ """
+
+def sparse(
+ sparsity: float, mean: float = ..., std: float = ..., dtype: mx.Dtype = ...
+) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns a sparse matrix.
+
+ Args:
+ sparsity (float): The fraction of elements in each column to be set to
+ zero.
+ mean (float, optional): Mean of the normal distribution. Default:
+ ``0.0``.
+ std (float, optional): Standard deviation of the normal distribution.
+ Default: ``1.0``.
+ dtype (Dtype, optional): The data type of the array. Default:
+ ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns an array with the
+ same shape as the input, filled with samples from a normal distribution.
+
+ Example:
+
+ >>> init_fn = nn.init.sparse(sparsity=0.5)
+ >>> init_fn(mx.zeros((2, 2)))
+ array([[-1.91187, -0.117483],
+ [0, 0]], dtype=float32)
+ """
+
+def orthogonal(
+ gain: float = ..., dtype: mx.Dtype = ...
+) -> Callable[[mx.array], mx.array]:
+ r"""An initializer that returns an orthogonal matrix.
+
+ Args:
+ gain (float, optional): Scaling factor for the orthogonal matrix.
+ Default: ``1.0``.
+ dtype (Dtype, optional): Data type of the array. Default: ``float32``.
+
+ Returns:
+ Callable[[array], array]: An initializer that returns
+ an orthogonal matrix with the same shape as the input.
+ """
diff --git a/typings/mlx/nn/layers/__init__.pyi b/typings/mlx/nn/layers/__init__.pyi
new file mode 100644
index 00000000..f22856cd
--- /dev/null
+++ b/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 *
diff --git a/typings/mlx/nn/layers/activations.pyi b/typings/mlx/nn/layers/activations.pyi
new file mode 100644
index 00000000..adacb4da
--- /dev/null
+++ b/typings/mlx/nn/layers/activations.pyi
@@ -0,0 +1,523 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from functools import partial
+from typing import Any
+
+import mlx.core as mx
+from base import Module
+
+@partial(mx.compile, shapeless=True)
+def sigmoid(x: mx.array) -> mx.array:
+ r"""Applies the sigmoid function.
+
+ .. math::
+ \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
+ """
+
+@partial(mx.compile, shapeless=True)
+def relu(x: mx.array) -> mx.array:
+ r"""Applies the Rectified Linear Unit.
+
+ Simply ``mx.maximum(x, 0)``.
+ """
+
+@partial(mx.compile, shapeless=True)
+def relu2(x: mx.array) -> mx.array:
+ r"""Applies the ReLU² activation function.
+
+ Applies :math:`\max(0, x)^2` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def relu6(x: mx.array) -> mx.array:
+ r"""Applies the Rectified Linear Unit 6.
+
+ Applies :math:`\min(\max(x, 0), 6)` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def leaky_relu(x: mx.array, negative_slope=...) -> mx.array:
+ r"""Applies the Leaky Rectified Linear Unit.
+
+ Simply ``mx.maximum(negative_slope * x, x)``.
+ """
+
+@partial(mx.compile, shapeless=True)
+def log_softmax(x: mx.array, axis=...):
+ r"""Applies the Log Softmax function.
+
+ Applies :math:`x + \log \sum_i e^{x_i}` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def elu(x: mx.array, alpha=...) -> mx.array:
+ r"""Applies the Exponential Linear Unit.
+
+ Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``.
+ """
+
+@partial(mx.compile, shapeless=True)
+def softmax(x: mx.array, axis=...) -> mx.array:
+ r"""Applies the Softmax function.
+
+ Applies :math:`\frac{e^{x_i}}{\sum_j e^{x_j}}` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def softplus(x: mx.array) -> mx.array:
+ r"""Applies the Softplus function.
+
+ Applies :math:`\log(1 + \exp(x))` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def softsign(x: mx.array) -> mx.array:
+ r"""Applies the Softsign function.
+
+ Applies :math:`\frac{x}{1 + |x|}` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def softshrink(x: mx.array, lambd: float = ...) -> mx.array:
+ r"""Applies the Softshrink activation function.
+
+ .. math::
+ \text{softshrink}(x) = \begin{cases}
+ x - \lambda & \text{if } x > \lambda \\
+ x + \lambda & \text{if } x < -\lambda \\
+ 0 & \text{otherwise}
+ \end{cases}
+ """
+
+@partial(mx.compile, shapeless=True)
+def celu(x: mx.array, alpha=...) -> mx.array:
+ r"""Applies the Continuously Differentiable Exponential Linear Unit.
+
+ Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))`
+ element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def silu(x: mx.array) -> mx.array:
+ r"""Applies the Sigmoid Linear Unit. Also known as Swish.
+
+ Applies :math:`x \sigma(x)` element wise, where :math:`\sigma(\cdot)` is
+ the logistic sigmoid.
+ """
+
+@partial(mx.compile, shapeless=True)
+def log_sigmoid(x: mx.array) -> mx.array:
+ r"""Applies the Log Sigmoid function.
+
+ Applies :math:`\log(\sigma(x)) = -\log(1 + e^{-x})` element wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def gelu(x: mx.array) -> mx.array:
+ r"""Applies the Gaussian Error Linear Units function.
+
+ .. math::
+ \textrm{GELU}(x) = x * \Phi(x)
+
+ where :math:`\Phi(x)` is the Gaussian CDF.
+
+ See also :func:`gelu_approx` and :func:`gelu_fast_approx` for faster
+ approximations.
+ """
+
+@partial(mx.compile, shapeless=True)
+def gelu_approx(x: mx.array) -> mx.array:
+ r"""An approximation to Gaussian Error Linear Unit.
+
+ See :func:`gelu` for the exact computation.
+
+ This function approximates ``gelu`` with a maximum absolute error :math:`<
+ 0.0005` in the range :math:`[-6, 6]` using the following
+
+ .. math::
+
+ x = 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right)
+
+ """
+
+@partial(mx.compile, shapeless=True)
+def gelu_fast_approx(x: mx.array) -> mx.array:
+ r"""A fast approximation to Gaussian Error Linear Unit.
+
+ See :func:`gelu` for the exact computation.
+
+ This function approximates ``gelu`` with a maximum absolute error :math:`<
+ 0.015` in the range :math:`[-6, 6]` using the following
+
+ .. math::
+
+ x = x \sigma\left(1.702 x\right)
+
+ where :math:`\sigma(\cdot)` is the logistic sigmoid.
+
+ References:
+ - https://github.com/hendrycks/GELUs
+ - https://arxiv.org/abs/1606.08415
+ """
+
+def glu(x: mx.array, axis: int = ...) -> mx.array:
+ r"""Applies the gated linear unit function.
+
+ This function splits the ``axis`` dimension of the input into two halves
+ (:math:`a` and :math:`b`) and applies :math:`a * \sigma(b)`.
+
+ .. math::
+ \textrm{GLU}(x) = a * \sigma(b)
+
+ Args:
+ axis (int): The dimension to split along. Default: ``-1``
+ """
+
+@partial(mx.compile, shapeless=True)
+def step(x: mx.array, threshold: float = ...) -> mx.array:
+ r"""Applies the Step Activation Function.
+
+ This function implements a binary step activation, where the output is set
+ to 1 if the input is greater than a specified threshold, and 0 otherwise.
+
+ .. math::
+ \text{step}(x) = \begin{cases}
+ 0 & \text{if } x < \text{threshold} \\
+ 1 & \text{if } x \geq \text{threshold}
+ \end{cases}
+
+ Args:
+ threshold: The value to threshold at.
+ """
+
+@partial(mx.compile, shapeless=True)
+def selu(x: mx.array) -> mx.array:
+ r"""Applies the Scaled Exponential Linear Unit.
+
+ .. math::
+ \text{selu}(x) = \begin{cases}
+ \lambda x & \text{if } x > 0 \\
+ \lambda \alpha (\exp(x) - 1) & \text{if } x \leq 0
+ \end{cases}
+
+ where :math:`\lambda = 1.0507` and :math:`\alpha = 1.67326`.
+
+ See also :func:`elu`.
+ """
+
+@partial(mx.compile, shapeless=True)
+def prelu(x: mx.array, alpha: mx.array) -> mx.array:
+ r"""Applies the element-wise parametric ReLU.
+
+ .. math::
+ \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
+
+ where :math:`a` is an array.
+ """
+
+@partial(mx.compile, shapeless=True)
+def mish(x: mx.array) -> mx.array:
+ r"""Applies the Mish function, element-wise.
+
+ Mish: A Self Regularized Non-Monotonic Neural Activation Function.
+
+ Reference: https://arxiv.org/abs/1908.08681
+
+ .. math::
+ \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
+
+ """
+
+@partial(mx.compile, shapeless=True)
+def hardswish(x: mx.array) -> mx.array:
+ r"""Applies the hardswish function, element-wise.
+
+ .. math::
+ \text{Hardswish}(x) = x * \min(\max(x + 3, 0), 6) / 6
+ """
+
+@partial(mx.compile, shapeless=True)
+def hard_tanh(x: mx.array, min_val=..., max_val=...) -> mx.array:
+ r"""Applies the HardTanh function.
+
+ Applies :math:`\max(\min(x, \text{max\_val}), \text{min\_val})` element-wise.
+ """
+
+@partial(mx.compile, shapeless=True)
+def hard_shrink(x: mx.array, lambd=...) -> mx.array:
+ r"""Applies the HardShrink activation function.
+
+ .. math::
+ \text{hardshrink}(x) = \begin{cases}
+ x & \text{if } x > \lambda \\
+ x & \text{if } x < -\lambda \\
+ 0 & \text{otherwise}
+ \end{cases}
+ """
+
+@partial(mx.compile, shapeless=True)
+def softmin(x: mx.array, axis=...) -> mx.array:
+ r"""Applies the Softmin function.
+
+ Applies :math:`\frac{e^{-x_i}}{\sum_j e^{-x_j}}` element-wise.
+ """
+
+def tanh(x: mx.array) -> mx.array:
+ """Applies the hyperbolic tangent function.
+
+ Simply ``mx.tanh(x)``.
+ """
+
+class GLU(Module):
+ r"""Applies the gated linear unit function.
+
+ This function splits the ``axis`` dimension of the input into two halves
+ (:math:`a` and :math:`b`) and applies :math:`a * \sigma(b)`.
+
+ .. math::
+ \textrm{GLU}(x) = a * \sigma(b)
+
+ Args:
+ axis (int): The dimension to split along. Default: ``-1``
+ """
+ def __init__(self, axis: int = ...) -> None: ...
+ def __call__(self, x) -> Any: ...
+
+@_make_activation_module(sigmoid)
+class Sigmoid(Module):
+ r"""Applies the sigmoid function, element-wise.
+
+ .. math::
+ \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
+ """
+
+@_make_activation_module(mish)
+class Mish(Module):
+ r"""Applies the Mish function, element-wise.
+
+ Reference: https://arxiv.org/abs/1908.08681
+
+ .. math::
+ \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
+
+ """
+
+@_make_activation_module(relu)
+class ReLU(Module):
+ r"""Applies the Rectified Linear Unit.
+ Simply ``mx.maximum(x, 0)``.
+
+ See :func:`relu` for the functional equivalent.
+ """
+
+@_make_activation_module(relu2)
+class ReLU2(Module):
+ r"""Applies the ReLU² activation function.
+
+ See :func:`relu2` for the functional equivalent.
+ """
+
+@_make_activation_module(relu6)
+class ReLU6(Module):
+ r"""Applies the Rectified Linear Unit 6.
+
+ See :func:`relu6` for the functional equivalent.
+ """
+
+class LeakyReLU(Module):
+ r"""Applies the Leaky Rectified Linear Unit.
+
+ Simply ``mx.maximum(negative_slope * x, x)``.
+
+ Args:
+ negative_slope: Controls the angle of the negative slope. Default: ``1e-2``
+ """
+ def __init__(self, negative_slope=...) -> None: ...
+ def __call__(self, x): ...
+
+class ELU(Module):
+ r"""Applies the Exponential Linear Unit.
+ Simply ``mx.where(x > 0, x, alpha * (mx.exp(x) - 1))``.
+
+ See :func:`elu` for the functional equivalent.
+
+ Args:
+ alpha: the :math:`\alpha` value for the ELU formulation. Default: ``1.0``
+ """
+ def __init__(self, alpha=...) -> None: ...
+ def __call__(self, x): ...
+
+@_make_activation_module(softmax)
+class Softmax(Module):
+ r"""Applies the Softmax function.
+
+ See :func:`softmax` for the functional equivalent.
+ """
+
+@_make_activation_module(softplus)
+class Softplus(Module):
+ r"""Applies the Softplus function.
+
+ See :func:`softplus` for the functional equivalent.
+ """
+
+@_make_activation_module(softsign)
+class Softsign(Module):
+ r"""Applies the Softsign function.
+
+ See :func:`softsign` for the functional equivalent.
+ """
+
+class Softshrink(Module):
+ r"""Applies the Softshrink function.
+
+ See :func:`softshrink` for the functional equivalent.
+
+ Args:
+ lambd: the :math:`\lambda` value for Softshrink. Default: ``0.5``
+ """
+ def __init__(self, lambd=...) -> None: ...
+ def __call__(self, x): ...
+
+class CELU(Module):
+ r"""Applies the Continuously Differentiable Exponential Linear Unit.
+ Applies :math:`\max(0, x) + \min(0, \alpha * (\exp(x / \alpha) - 1))`
+ element wise.
+
+ See :func:`celu` for the functional equivalent.
+
+ Args:
+ alpha: the :math:`\alpha` value for the CELU formulation. Default: ``1.0``
+ """
+ def __init__(self, alpha=...) -> None: ...
+ def __call__(self, x): ...
+
+@_make_activation_module(silu)
+class SiLU(Module):
+ r"""Applies the Sigmoid Linear Unit. Also known as Swish.
+
+ See :func:`silu` for the functional equivalent.
+ """
+
+@_make_activation_module(log_softmax)
+class LogSoftmax(Module):
+ r"""Applies the Log Softmax function.
+
+ See :func:`log_softmax` for the functional equivalent.
+ """
+
+@_make_activation_module(log_sigmoid)
+class LogSigmoid(Module):
+ r"""Applies the Log Sigmoid function.
+
+ See :func:`log_sigmoid` for the functional equivalent.
+ """
+
+class PReLU(Module):
+ r"""Applies the element-wise parametric ReLU.
+ Applies :math:`\max(0, x) + a * \min(0, x)` element wise, where :math:`a`
+ is an array.
+
+ See :func:`prelu` for the functional equivalent.
+
+ Args:
+ num_parameters: number of :math:`a` to learn. Default: ``1``
+ init: the initial value of :math:`a`. Default: ``0.25``
+ """
+ def __init__(self, num_parameters=..., init=...) -> None: ...
+ def __call__(self, x: mx.array): ...
+
+class GELU(Module):
+ r"""Applies the Gaussian Error Linear Units.
+
+ .. math::
+ \textrm{GELU}(x) = x * \Phi(x)
+
+ where :math:`\Phi(x)` is the Gaussian CDF.
+
+ However, if ``approx`` is set to 'precise' or 'fast' it applies
+
+ .. math::
+ \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) \\
+ \textrm{GELUFast}(x) &= x * \sigma\left(1.702 * x\right)
+
+ respectively.
+
+ .. note::
+ For compatibility with the PyTorch API, 'tanh' can be used as an alias
+ for 'precise'.
+
+ See :func:`gelu`, :func:`gelu_approx` and :func:`gelu_fast_approx` for the
+ functional equivalents and information regarding error bounds.
+
+
+ Args:
+ approx ('none' | 'precise' | 'fast'): Which approximation to gelu to use if any.
+ """
+ def __init__(self, approx=...) -> None: ...
+ def __call__(self, x): ...
+
+@_make_activation_module(tanh)
+class Tanh(Module):
+ r"""Applies the hyperbolic tangent function.
+
+ See :func:`tanh` for the functional equivalent.
+ """
+
+@_make_activation_module(hardswish)
+class Hardswish(Module):
+ r"""Applies the hardswish function, element-wise.
+
+ See :func:`hardswish` for the functional equivalent.
+ """
+
+class Step(Module):
+ r"""Applies the Step Activation Function.
+
+ This function implements a binary step activation, where the output is set
+ to 1 if the input is greater than a specified threshold, and 0 otherwise.
+
+ .. math::
+ \text{step}(x) = \begin{cases}
+ 0 & \text{if } x < \text{threshold} \\
+ 1 & \text{if } x \geq \text{threshold}
+ \end{cases}
+
+ Args:
+ threshold: The value to threshold at.
+ """
+ def __init__(self, threshold: float = ...) -> None: ...
+ def __call__(self, x: mx.array): ...
+
+@_make_activation_module(selu)
+class SELU(Module):
+ r"""Applies the Scaled Exponential Linear Unit.
+
+ See :func:`selu` for the functional equivalent.
+ """
+
+@_make_activation_module(hard_tanh)
+class HardTanh(Module):
+ r"""Applies the HardTanh function.
+
+ See :func:`hard_tanh` for the functional equivalent.
+ """
+
+@_make_activation_module(hard_shrink)
+class HardShrink(Module):
+ r"""Applies the HardShrink function.
+
+ See :func:`hard_shrink` for the functional equivalent.
+
+ Args:
+ lambd: the :math:`\lambda` value for Hardshrink. Default: ``0.5``
+ """
+
+@_make_activation_module(softmin)
+class Softmin(Module):
+ r"""Applies the Softmin function.
+
+ See :func:`softmin` for the functional equivalent.
+ """
diff --git a/typings/mlx/nn/layers/base.pyi b/typings/mlx/nn/layers/base.pyi
new file mode 100644
index 00000000..a4abf36b
--- /dev/null
+++ b/typings/mlx/nn/layers/base.pyi
@@ -0,0 +1,393 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import mlx.core as mx
+
+class Module(dict):
+ """Base class for building neural networks with MLX.
+
+ All the layers provided in :mod:`layers` subclass this class and
+ your models should do the same.
+
+ A ``Module`` can contain other ``Module`` instances or :class:`mlx.core.array`
+ instances in arbitrary nesting of python lists or dicts. The ``Module``
+ then allows recursively extracting all the :class:`mlx.core.array` instances
+ using :meth:`Module.parameters`.
+
+ In addition, the ``Module`` has the concept of trainable and non trainable
+ parameters (called "frozen"). When using :func:`value_and_grad`
+ the gradients are returned only with respect to the trainable parameters.
+ All arrays in a module are trainable unless they are added in the "frozen"
+ set by calling :meth:`freeze`.
+
+ .. code-block:: python
+
+ import mlx.core as mx
+ import mlx.nn as nn
+
+ class MyMLP(nn.Module):
+ def __init__(self, in_dims: int, out_dims: int, hidden_dims: int = 16):
+ super().__init__()
+
+ self.in_proj = nn.Linear(in_dims, hidden_dims)
+ self.out_proj = nn.Linear(hidden_dims, out_dims)
+
+ def __call__(self, x):
+ x = self.in_proj(x)
+ x = mx.maximum(x, 0)
+ return self.out_proj(x)
+
+ model = MyMLP(2, 1)
+
+ # All the model parameters are created but since MLX is lazy by
+ # default, they are not evaluated yet. Calling `mx.eval` actually
+ # allocates memory and initializes the parameters.
+ mx.eval(model.parameters())
+
+ # Setting a parameter to a new value is as simply as accessing that
+ # parameter and assigning a new array to it.
+ model.in_proj.weight = model.in_proj.weight * 2
+ mx.eval(model.parameters())
+ """
+
+ __call__: Callable
+ def __init__(self) -> None:
+ """Should be called by the subclasses of ``Module``."""
+
+ @property
+ def training(self): # -> bool:
+ """Boolean indicating if the model is in training mode."""
+
+ @property
+ def state(self): # -> Self:
+ """The module's state dictionary
+
+ The module's state dictionary contains any attribute set on the
+ module including parameters in :meth:`Module.parameters`
+
+ Unlike :meth:`Module.parameters`, the :attr:`Module.state` property is
+ a reference to the module's state. Updates to it will be reflected in
+ the original module.
+ """
+
+ def __repr__(self): # -> str:
+ ...
+ def __getattr__(self, key: str): # -> None:
+ ...
+ def __setattr__(self, key: str, val: Any): # -> None:
+ ...
+ def __delattr__(self, name): # -> None:
+ ...
+ def load_weights(
+ self,
+ file_or_weights: Union[str, List[Tuple[str, mx.array]]],
+ strict: bool = ...,
+ ) -> Module:
+ """
+ Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list.
+
+ Args:
+ file_or_weights (str or list(tuple(str, mx.array))): The path to
+ the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list
+ of pairs of parameter names and arrays.
+ strict (bool, optional): If ``True`` then checks that the provided
+ weights exactly match the parameters of the model. Otherwise,
+ only the weights actually contained in the model are loaded and
+ shapes are not checked. Default: ``True``.
+
+ Returns:
+ The module instance after updating the weights.
+
+ Example:
+
+ .. code-block:: python
+
+ import mlx.core as mx
+ import mlx.nn as nn
+ model = nn.Linear(10, 10)
+
+ # Load from file
+ model.load_weights("weights.npz")
+
+ # Load from .safetensors file
+ model.load_weights("weights.safetensors")
+
+ # Load from list
+ weights = [
+ ("weight", mx.random.uniform(shape=(10, 10))),
+ ("bias", mx.zeros((10,))),
+ ]
+ model.load_weights(weights)
+
+ # Missing weight
+ weights = [
+ ("weight", mx.random.uniform(shape=(10, 10))),
+ ]
+
+ # Raises a ValueError exception
+ model.load_weights(weights)
+
+ # Ok, only updates the weight but not the bias
+ model.load_weights(weights, strict=False)
+ """
+
+ def save_weights(self, file: str): # -> None:
+ """
+ Save the model's weights to a file. The saving method is determined by the file extension:
+ - ``.npz`` will use :func:`mx.savez`
+ - ``.safetensors`` will use :func:`mx.save_safetensors`
+ """
+
+ @staticmethod
+ def is_module(value): # -> bool:
+ ...
+ @staticmethod
+ def valid_child_filter(module, key, value): # -> bool:
+ ...
+ @staticmethod
+ def valid_parameter_filter(module, key, value): # -> bool:
+ ...
+ @staticmethod
+ def trainable_parameter_filter(module, key, value): # -> bool:
+ ...
+ def filter_and_map(
+ self,
+ filter_fn: Callable[[Module, str, Any], bool],
+ map_fn: Optional[Callable] = ...,
+ is_leaf_fn: Optional[Callable[[Module, str, Any], bool]] = ...,
+ ): # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
+ """Recursively filter the contents of the module using ``filter_fn``,
+ namely only select keys and values where ``filter_fn`` returns true.
+
+ This is used to implement :meth:`parameters` and :meth:`trainable_parameters`
+ but it can also be used to extract any subset of the module's parameters.
+
+ Args:
+ filter_fn (Callable): Given a value, the key in which it is found
+ and the containing module, decide whether to keep the value or
+ drop it.
+ map_fn (Callable, optional): Optionally transform the value before
+ returning it.
+ is_leaf_fn (Callable, optional): Given a value, the key in which it
+ is found and the containing module decide if it is a leaf.
+
+ Returns:
+ A dictionary containing the contents of the module recursively filtered
+ """
+
+ def parameters(
+ self,
+ ) -> mx.MX_ARRAY_TREE:
+ """Recursively return all the :class:`mlx.core.array` members of this Module
+ as a dict of dicts and lists."""
+
+ def trainable_parameters(
+ self,
+ ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
+ """Recursively return all the non frozen :class:`mlx.core.array` members of
+ this Module as a dict of dicts and lists."""
+
+ def children(
+ self,
+ ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
+ """Return the direct descendants of this Module instance."""
+
+ def leaf_modules(
+ self,
+ ) -> 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, strict: bool = ...) -> Module:
+ """Replace the parameters of this Module with the provided ones in the
+ dict of dicts and lists.
+
+ Commonly used by the optimizer to change the model to the updated
+ (optimized) parameters. Also used by the :meth:`value_and_grad` to set the
+ tracers in the model in order to compute gradients.
+
+ The passed in parameters dictionary need not be a full dictionary
+ similar to :meth:`parameters`. Only the provided locations will be
+ updated.
+
+ Args:
+ parameters (dict): A complete or partial dictionary of the modules
+ parameters.
+ strict (bool): If ``True`` checks that ``parameters`` is a
+ subset of the module's parameters. Default: ``True``.
+ Returns:
+ The module instance after updating the parameters.
+ """
+
+ def apply(
+ self,
+ map_fn: Callable[[mx.array], mx.array],
+ filter_fn: Optional[Callable[[Module, str, Any], bool]] = ...,
+ ) -> Module:
+ """Map all the parameters using the provided ``map_fn`` and immediately
+ update the module with the mapped parameters.
+
+ For instance running ``model.apply(lambda x: x.astype(mx.float16))``
+ casts all parameters to 16 bit floats.
+
+ Args:
+ map_fn (Callable): Maps an array to another array
+ filter_fn (Callable, optional): Filter to select which arrays to
+ map (default: :meth:`Module.valid_parameter_filter`).
+
+ Returns:
+ The module instance after updating the parameters.
+ """
+
+ def update_modules(self, modules: dict, strict: bool = ...) -> Module:
+ """Replace the child modules of this :class:`Module` instance with the
+ provided ones in the dict of dicts and lists.
+
+ It is the equivalent of :meth:`Module.update` but for modules instead
+ of parameters and allows us to flexibly edit complex architectures by
+ programmatically swapping layers.
+
+ The passed in parameters dictionary need not be a full dictionary
+ similar to :meth:`modules`. Only the provided locations will be
+ updated.
+
+ Args:
+ modules (dict): A complete or partial dictionary of the module's
+ submodules.
+ strict (bool): If ``True`` checks that ``modules`` is a
+ subset of the child modules of this instance. Default: ``True``.
+ Returns:
+ The module instance after updating the submodules.
+ """
+
+ def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module:
+ """Apply a function to all the modules in this instance (including this
+ instance).
+
+ Args:
+ apply_fn (Callable): The function to apply to the modules.
+
+ Returns:
+ The module instance after updating submodules.
+ """
+
+ def modules(self): # -> list[Any]:
+ """Return a list with all the modules in this instance.
+
+ Returns:
+ A list of :class:`Module` instances.
+ """
+
+ def named_modules(self): # -> list[Any]:
+ """Return a list with all the modules in this instance and their name
+ with dot notation.
+
+ Returns:
+ A list of tuples (str, :class:`Module`).
+ """
+
+ def freeze(
+ self,
+ *,
+ recurse: bool = ...,
+ keys: Optional[Union[str, List[str]]] = ...,
+ strict: bool = ...,
+ ) -> Module:
+ """Freeze the Module's parameters or some of them. Freezing a parameter means not
+ computing gradients for it.
+
+ This function is idempotent i.e. freezing a frozen model is a no-op.
+
+ Example:
+ For instance to only train the attention parameters from a Transformer:
+
+ .. code-block:: python
+
+ model = nn.Transformer()
+ model.freeze()
+ model.apply_to_modules(lambda k, v: v.unfreeze() if k.endswith("attention") else None)
+
+ Args:
+ recurse (bool, optional): If True then freeze the parameters of the
+ submodules as well. Default: ``True``.
+ keys (str or list[str], optional): If provided then only these
+ parameters will be frozen otherwise all the parameters of a
+ module. For instance freeze all biases by calling
+ ``module.freeze(keys="bias")``.
+ strict (bool, optional): If set to ``True`` validate that the passed keys exist.
+ Default: ``False``.
+
+ Returns:
+ The module instance after freezing the parameters.
+ """
+
+ def unfreeze(
+ self,
+ *,
+ recurse: bool = ...,
+ keys: Optional[Union[str, List[str]]] = ...,
+ strict: bool = ...,
+ ) -> Module:
+ """Unfreeze the Module's parameters or some of them.
+
+ This function is idempotent ie unfreezing a model that is not frozen is
+ a noop.
+
+ Example:
+
+ For instance to only train the biases of a Transformer one can do:
+
+ .. code-block:: python
+
+ model = nn.Transformer()
+ model.freeze()
+ model.unfreeze(keys="bias")
+
+ Args:
+ recurse (bool, optional): If True then unfreeze the parameters of the
+ submodules as well. Default: ``True``.
+ keys (str or list[str], optional): If provided then only these
+ parameters will be unfrozen otherwise all the parameters of a
+ module. For instance unfreeze all biases by calling
+ ``module.unfreeze(keys="bias")``.
+ strict (bool, optional): If set to ``True`` validate that the passed keys exist.
+ Default: ``False``.
+
+ Returns:
+ The module instance after unfreezing the parameters.
+ """
+
+ def train(self, mode: bool = ...) -> Module:
+ """Set the model in or out of training mode.
+
+ Training mode only applies to certain layers. For example
+ :obj:`Dropout` applies a random mask in training mode, but is the
+ identity in evaluation mode.
+
+ Args:
+ mode (bool): Indicate if the model should be in training or
+ evaluation mode. Default: ``True``.
+ Returns:
+ The module instance after updating the training mode.
+ """
+
+ def eval(self) -> Module:
+ """Set the model to evaluation mode.
+
+ See :func:`train`.
+ """
+
+ def set_dtype(
+ self, dtype: mx.Dtype, predicate: Optional[Callable[[mx.Dtype], bool]] = ...
+ ): # -> None:
+ """Set the dtype of the module's parameters.
+
+ Args:
+ dtype (Dtype): The new dtype.
+ predicate (typing.Callable, optional): A predicate to select
+ parameters to cast. By default, only parameters of type
+ :attr:`floating` will be updated to avoid casting integer
+ parameters to the new dtype.
+ """
diff --git a/typings/mlx/nn/layers/containers.pyi b/typings/mlx/nn/layers/containers.pyi
new file mode 100644
index 00000000..068ea179
--- /dev/null
+++ b/typings/mlx/nn/layers/containers.pyi
@@ -0,0 +1,21 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Callable
+
+import mlx.core as mx
+from base import Module
+
+class Sequential(Module):
+ """A layer that calls the passed callables in order.
+
+ We can pass either modules or plain callables to the Sequential module. If
+ our functions have learnable parameters they should be implemented as
+ ``nn.Module`` instances.
+
+ Args:
+ modules (tuple of Callables): The modules to call in order
+ """
+ def __init__(self, *modules: Module | Callable[[mx.array], mx.array]) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/convolution.pyi b/typings/mlx/nn/layers/convolution.pyi
new file mode 100644
index 00000000..c68ad289
--- /dev/null
+++ b/typings/mlx/nn/layers/convolution.pyi
@@ -0,0 +1,116 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Union
+
+import mlx.core as mx
+from base import Module
+
+class Conv1d(Module):
+ """Applies a 1-dimensional convolution over the multi-channel input sequence.
+
+ The channels are expected to be last i.e. the input shape should be ``NLC`` where:
+
+ * ``N`` is the batch dimension
+ * ``L`` is the sequence length
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels
+ out_channels (int): The number of output channels
+ kernel_size (int): The size of the convolution filters
+ stride (int, optional): The stride when applying the filter.
+ Default: ``1``.
+ padding (int, optional): How many positions to 0-pad the input with.
+ Default: ``0``.
+ dilation (int, optional): The dilation of the convolution.
+ groups (int, optional): The number of groups for the convolution.
+ Default: ``1``.
+ bias (bool, optional): If ``True`` add a learnable bias to the output.
+ Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = ...,
+ padding: int = ...,
+ dilation: int = ...,
+ groups: int = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class Conv2d(Module):
+ """Applies a 2-dimensional convolution over the multi-channel input image.
+
+ The channels are expected to be last i.e. the input shape should be ``NHWC`` where:
+
+ * ``N`` is the batch dimension
+ * ``H`` is the input image height
+ * ``W`` is the input image width
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels.
+ out_channels (int): The number of output channels.
+ kernel_size (int or tuple): The size of the convolution filters.
+ stride (int or tuple, optional): The size of the stride when
+ applying the filter. Default: ``1``.
+ padding (int or tuple, optional): How many positions to 0-pad
+ the input with. Default: ``0``.
+ dilation (int or tuple, optional): The dilation of the convolution.
+ groups (int, optional): The number of groups for the convolution.
+ Default: ``1``.
+ bias (bool, optional): If ``True`` add a learnable bias to the
+ output. Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, tuple],
+ stride: Union[int, tuple] = ...,
+ padding: Union[int, tuple] = ...,
+ dilation: Union[int, tuple] = ...,
+ groups: int = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x) -> mx.array: ...
+
+class Conv3d(Module):
+ """Applies a 3-dimensional convolution over the multi-channel input image.
+
+ The channels are expected to be last i.e. the input shape should be ``NDHWC`` where:
+
+ * ``N`` is the batch dimension
+ * ``D`` is the input image depth
+ * ``H`` is the input image height
+ * ``W`` is the input image width
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels.
+ out_channels (int): The number of output channels.
+ kernel_size (int or tuple): The size of the convolution filters.
+ stride (int or tuple, optional): The size of the stride when
+ applying the filter. Default: ``1``.
+ dilation (int or tuple, optional): The dilation of the convolution.
+ padding (int or tuple, optional): How many positions to 0-pad
+ the input with. Default: ``0``.
+ bias (bool, optional): If ``True`` add a learnable bias to the
+ output. Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, tuple],
+ stride: Union[int, tuple] = ...,
+ padding: Union[int, tuple] = ...,
+ dilation: Union[int, tuple] = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/convolution_transpose.pyi b/typings/mlx/nn/layers/convolution_transpose.pyi
new file mode 100644
index 00000000..8fe11b4a
--- /dev/null
+++ b/typings/mlx/nn/layers/convolution_transpose.pyi
@@ -0,0 +1,119 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Union
+
+import mlx.core as mx
+from base import Module
+
+class ConvTranspose1d(Module):
+ """Applies a 1-dimensional transposed convolution over the multi-channel input sequence.
+
+ The channels are expected to be last i.e. the input shape should be ``NLC`` where:
+
+ * ``N`` is the batch dimension
+ * ``L`` is the sequence length
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels
+ out_channels (int): The number of output channels
+ kernel_size (int): The size of the convolution filters
+ stride (int, optional): The stride when applying the filter.
+ Default: ``1``.
+ padding (int, optional): How many positions to 0-pad the input with.
+ Default: ``0``.
+ dilation (int, optional): The dilation of the convolution.
+ output_padding(int, optional): Additional size added to one side of the
+ output shape. Default: ``0``.
+ bias (bool, optional): If ``True`` add a learnable bias to the output.
+ Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = ...,
+ padding: int = ...,
+ dilation: int = ...,
+ output_padding: int = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class ConvTranspose2d(Module):
+ """Applies a 2-dimensional transposed convolution over the multi-channel input image.
+
+ The channels are expected to be last i.e. the input shape should be ``NHWC`` where:
+
+ * ``N`` is the batch dimension
+ * ``H`` is the input image height
+ * ``W`` is the input image width
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels.
+ out_channels (int): The number of output channels.
+ kernel_size (int or tuple): The size of the convolution filters.
+ stride (int or tuple, optional): The size of the stride when
+ applying the filter. Default: ``1``.
+ padding (int or tuple, optional): How many positions to 0-pad
+ the input with. Default: ``0``.
+ dilation (int or tuple, optional): The dilation of the convolution.
+ output_padding(int or tuple, optional): Additional size added to one
+ side of the output shape. Default: ``0``.
+ bias (bool, optional): If ``True`` add a learnable bias to the
+ output. Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, tuple],
+ stride: Union[int, tuple] = ...,
+ padding: Union[int, tuple] = ...,
+ dilation: Union[int, tuple] = ...,
+ output_padding: Union[int, tuple] = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class ConvTranspose3d(Module):
+ """Applies a 3-dimensional transposed convolution over the multi-channel input image.
+
+ The channels are expected to be last i.e. the input shape should be ``NDHWC`` where:
+
+ * ``N`` is the batch dimension
+ * ``D`` is the input image depth
+ * ``H`` is the input image height
+ * ``W`` is the input image width
+ * ``C`` is the number of input channels
+
+ Args:
+ in_channels (int): The number of input channels.
+ out_channels (int): The number of output channels.
+ kernel_size (int or tuple): The size of the convolution filters.
+ stride (int or tuple, optional): The size of the stride when
+ applying the filter. Default: ``1``.
+ padding (int or tuple, optional): How many positions to 0-pad
+ the input with. Default: ``0``.
+ dilation (int or tuple, optional): The dilation of the convolution.
+ output_padding(int or tuple, optional): Additional size added to one
+ side of the output shape. Default: ``0``.
+ bias (bool, optional): If ``True`` add a learnable bias to the
+ output. Default: ``True``
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, tuple],
+ stride: Union[int, tuple] = ...,
+ padding: Union[int, tuple] = ...,
+ dilation: Union[int, tuple] = ...,
+ output_padding: Union[int, tuple] = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/distributed.pyi b/typings/mlx/nn/layers/distributed.pyi
new file mode 100644
index 00000000..5be9cc4b
--- /dev/null
+++ b/typings/mlx/nn/layers/distributed.pyi
@@ -0,0 +1,227 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from functools import lru_cache
+from typing import Callable, Optional, Union
+
+import mlx.core as mx
+from base import Module
+from mlx.nn.layers.linear import Linear
+
+@lru_cache
+def sum_gradients(
+ group: mx.distributed.Group,
+) -> Callable[..., mx.array]: # -> Callable[..., Any] | Callable[..., array]:
+ ...
+def shard_inplace(
+ module: Module,
+ sharding: str,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+) -> None:
+ """Shard a module in-place by updating its parameter dictionary with the
+ sharded parameter dictionary.
+
+ The ``sharding`` argument can be any callable that given the path and the
+ weight returns the sharding axis and optionally also the segments that
+ comprise the unsharded weight. For instance if the weight is a fused QKV
+ matrix the segments should be 3.
+
+ .. note::
+ The module doesn't change so in order for distributed communication to
+ happen the module needs to natively support it and for it to be enabled.
+
+ Args:
+ module (Module): The parameters of this module will be sharded
+ in-place.
+ sharding (str or callable): One of "all-to-sharded" and
+ "sharded-to-all" or a callable that returns the sharding axis and
+ segments.
+ segments (int or list): The segments to use if ``sharding`` is a
+ string. Default: ``1``.
+ group (mlx.core.distributed.Group): The distributed group to shard
+ across. If not set, the global group will be used. Default: ``None``.
+ """
+
+def shard_linear(
+ module: Module,
+ sharding: str,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+) -> Linear:
+ """Create a new linear layer that has its parameters sharded and also
+ performs distributed communication either in the forward or backward
+ pass.
+
+ .. note::
+ Contrary to ``shard_inplace``, the original layer is not changed but a
+ new layer is returned.
+
+ Args:
+ module (Module): The linear layer to be sharded.
+ sharding (str): One of "all-to-sharded" and
+ "sharded-to-all" that defines the type of sharding to perform.
+ segments (int or list): The segments to use. Default: ``1``.
+ group (mlx.core.distributed.Group): The distributed group to shard
+ across. If not set, the global group will be used. Default: ``None``.
+ """
+
+class AllToShardedLinear(Module):
+ """Each member of the group applies part of the affine transformation such
+ that the result is sharded across the group.
+
+ The gradients are automatically aggregated from each member of the group.
+
+ Args:
+ input_dims (int): The dimensionality of the input features
+ output_dims (int): The dimensionality of the output features
+ bias (bool, optional): If set to ``False`` the the layer will not use a
+ bias. Default is ``True``.
+ group (mx.distributed.Group, optional): The sharding will happen across
+ this group. If not set then the global group is used. Default is
+ ``None``.
+ """
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ bias: bool = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ @classmethod
+ def from_linear(
+ cls,
+ linear_layer: Module,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> AllToShardedLinear: ...
+
+class ShardedToAllLinear(Module):
+ """Each member of the group applies part of the affine transformation and
+ then aggregates the results.
+
+ All nodes will have the same exact result after this layer.
+
+ :class:`ShardedToAllLinear` provides a classmethod :meth:`from_linear` to
+ convert linear layers to sharded :obj:`ShardedToAllLinear` layers.
+
+ Args:
+ input_dims (int): The dimensionality of the input features
+ output_dims (int): The dimensionality of the output features
+ bias (bool, optional): If set to ``False`` the the layer will not use a
+ bias. Default is ``True``.
+ group (mx.distributed.Group, optional): The sharding will happen across
+ this group. If not set then the global group is used. Default is
+ ``None``.
+ """
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ bias: bool = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ @classmethod
+ def from_linear(
+ cls,
+ linear_layer: Module,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> ShardedToAllLinear: ...
+
+class QuantizedAllToShardedLinear(Module):
+ """Each member of the group applies part of the affine transformation with
+ a quantized matrix such that the result is sharded across the group.
+
+ It is the quantized equivalent of :class:`AllToShardedLinear`.
+ Similar to :class:`QuantizedLinear` its parameters are frozen and
+ will not be included in any gradient computation.
+
+ Args:
+ input_dims (int): The dimensionality of the input features.
+ output_dims (int): The dimensionality of the output features.
+ bias (bool, optional): If set to ``False`` then the layer will not use
+ a bias. Default: ``True``.
+ group_size (int, optional): The group size to use for the quantized
+ weight. See :func:`~mlx.core.quantize`. Default: ``64``.
+ bits (int, optional): The bit width to use for the quantized weight.
+ See :func:`~mlx.core.quantize`. Default: ``4``.
+ group (mx.distributed.Group, optional): The sharding will happen across
+ this group. If not set then the global group is used. Default is
+ ``None``.
+ """
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ bias: bool = ...,
+ group_size: int = ...,
+ bits: int = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> None: ...
+ def unfreeze(self, *args, **kwargs) -> None:
+ """Wrap unfreeze so that we unfreeze any layers we might contain but
+ our parameters will remain frozen."""
+
+ def __call__(self, x: mx.array) -> mx.array: ...
+ @classmethod
+ def from_quantized_linear(
+ cls,
+ quantized_linear_layer: Module,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> QuantizedAllToShardedLinear: ...
+
+class QuantizedShardedToAllLinear(Module):
+ """Each member of the group applies part of the affine transformation using
+ the quantized matrix and then aggregates the results.
+
+ All nodes will have the same exact result after this layer.
+
+ It is the quantized equivalent of :class:`ShardedToAllLinear`.
+ Similar to :class:`QuantizedLinear` its parameters are frozen and
+ will not be included in any gradient computation.
+
+ Args:
+ input_dims (int): The dimensionality of the input features.
+ output_dims (int): The dimensionality of the output features.
+ bias (bool, optional): If set to ``False`` then the layer will not use
+ a bias. Default: ``True``.
+ group_size (int, optional): The group size to use for the quantized
+ weight. See :func:`~mlx.core.quantize`. Default: ``64``.
+ bits (int, optional): The bit width to use for the quantized weight.
+ See :func:`~mlx.core.quantize`. Default: ``4``.
+ group (mx.distributed.Group, optional): The sharding will happen across
+ this group. If not set then the global group is used. Default is
+ ``None``.
+ """
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ bias: bool = ...,
+ group_size: int = ...,
+ bits: int = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> None: ...
+ def unfreeze(self, *args, **kwargs): # -> None:
+ """Wrap unfreeze so that we unfreeze any layers we might contain but
+ our parameters will remain frozen."""
+
+ def __call__(self, x: mx.array) -> mx.array: ...
+ @classmethod
+ def from_quantized_linear(
+ cls,
+ quantized_linear_layer: Module,
+ *,
+ segments: Union[int, list[int]] = ...,
+ group: Optional[mx.distributed.Group] = ...,
+ ) -> QuantizedShardedToAllLinear: ...
diff --git a/typings/mlx/nn/layers/dropout.pyi b/typings/mlx/nn/layers/dropout.pyi
new file mode 100644
index 00000000..00ef6f01
--- /dev/null
+++ b/typings/mlx/nn/layers/dropout.pyi
@@ -0,0 +1,65 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+from base import Module
+
+class Dropout(Module):
+ r"""Randomly zero a portion of the elements during training.
+
+ The remaining elements are multiplied with :math:`\frac{1}{1-p}` where
+ :math:`p` is the probability of zeroing an element. This is done so the
+ expected value of a given element will remain the same.
+
+ Args:
+ p (float): The probability to zero an element
+ """
+ def __init__(self, p: float = ...) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class Dropout2d(Module):
+ r"""Apply 2D channel-wise dropout during training.
+
+ Randomly zero out entire channels independently with probability :math:`p`.
+ This layer expects the channels to be last, i.e. the input shape should be
+ ``NWHC`` or ``WHC`` where:``N`` is the batch dimension,``H`` is the input
+ image height,``W`` is the input image width, and``C`` is the number of
+ input channels
+
+ The remaining channels are scaled by :math:`\frac{1}{1-p}` to
+ maintain the expected value of each element. Unlike traditional dropout,
+ which zeros individual entries, this layer zeros entire channels. This is
+ beneficial for early convolution layers where adjacent pixels are
+ correlated. In such case, traditional dropout may not effectively
+ regularize activations. For more details, see [1].
+
+ [1]: Thompson, J., Goroshin, R., Jain, A., LeCun, Y. and Bregler C., 2015.
+ Efficient Object Localization Using Convolutional Networks. CVPR 2015.
+
+ Args:
+ p (float): Probability of zeroing a channel during training.
+ """
+ def __init__(self, p: float = ...) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class Dropout3d(Module):
+ r"""Apply 3D channel-wise dropout during training.
+
+ Randomly zero out entire channels independently with probability :math:`p`.
+ This layer expects the channels to be last, i.e., the input shape should be
+ `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth,
+ `H` is the input image height, `W` is the input image width, and `C` is
+ the number of input channels.
+
+ The remaining channels are scaled by :math:`\frac{1}{1-p}` to
+ maintain the expected value of each element. Unlike traditional dropout,
+ which zeros individual entries, this layer zeros entire channels. This is
+ often beneficial for convolutional layers processing 3D data, like in
+ medical imaging or video processing.
+
+ Args:
+ p (float): Probability of zeroing a channel during training.
+ """
+ def __init__(self, p: float = ...) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/embedding.pyi b/typings/mlx/nn/layers/embedding.pyi
new file mode 100644
index 00000000..e273c801
--- /dev/null
+++ b/typings/mlx/nn/layers/embedding.pyi
@@ -0,0 +1,34 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+from base import Module
+
+from .quantized import QuantizedEmbedding
+
+class Embedding(Module):
+ """Implements a simple lookup table that maps each input integer to a
+ high-dimensional vector.
+
+ Typically used to embed discrete tokens for processing by neural networks.
+
+ Args:
+ num_embeddings (int): How many possible discrete tokens can we embed.
+ Usually called the vocabulary size.
+ dims (int): The dimensionality of the embeddings.
+ """
+ def __init__(self, num_embeddings: int, dims: int) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ def as_linear(self, x: mx.array) -> mx.array:
+ """
+ Call the embedding layer as a linear layer.
+
+ Use this for example when input embedding and output projection
+ weights are tied.
+ """
+
+ def to_quantized(
+ self, group_size: int = ..., bits: int = ..., mode: str = ...
+ ) -> QuantizedEmbedding:
+ """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer."""
diff --git a/typings/mlx/nn/layers/linear.pyi b/typings/mlx/nn/layers/linear.pyi
new file mode 100644
index 00000000..f9c91874
--- /dev/null
+++ b/typings/mlx/nn/layers/linear.pyi
@@ -0,0 +1,76 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Any
+
+import mlx.core as mx
+from base import Module
+
+from .quantized import QuantizedLinear
+
+class Identity(Module):
+ r"""A placeholder identity operator that is argument-insensitive.
+
+ Args:
+ args: any argument (unused)
+ kwargs: any keyword argument (unused)
+ """
+ def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class Linear(Module):
+ r"""Applies an affine transformation to the input.
+
+ Concretely:
+
+ .. math::
+
+ y = x W^\top + b
+
+ where:
+ where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``.
+
+ The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
+ where :math:`k = \frac{1}{\sqrt{D_i}}` and :math:`D_i` is equal to ``input_dims``.
+
+ Args:
+ input_dims (int): The dimensionality of the input features
+ output_dims (int): The dimensionality of the output features
+ bias (bool, optional): If set to ``False`` then the layer will
+ not use a bias. Default is ``True``.
+ """
+ def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ def to_quantized(
+ self, group_size: int = ..., bits: int = ..., mode: str = ...
+ ) -> QuantizedLinear:
+ """Return a :obj:`QuantizedLinear` layer that approximates this layer."""
+
+class Bilinear(Module):
+ r"""Applies a bilinear transformation to the inputs.
+
+ Concretely:
+
+ .. math::
+
+ y_i = x_1^\top W_i x_2 + b_i
+
+ where:
+ :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``,
+ and :math:`i` indexes the output dimension.
+
+ The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
+ where :math:`k = \frac{1}{\sqrt{D_1}}` and :math:`D_1` is ``input1_dims``.
+
+ Args:
+ input1_dims (int): The dimensionality of the input1 features
+ input2_dims (int): The dimensionality of the input2 features
+ output_dims (int): The dimensionality of the output features
+ bias (bool, optional): If set to ``False`` then the layer will
+ not use a bias. Default is ``True``.
+ """
+ def __init__(
+ self, input1_dims: int, input2_dims: int, output_dims: int, bias: bool = ...
+ ) -> None: ...
+ def __call__(self, x1: mx.array, x2: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/normalization.pyi b/typings/mlx/nn/layers/normalization.pyi
new file mode 100644
index 00000000..4116f860
--- /dev/null
+++ b/typings/mlx/nn/layers/normalization.pyi
@@ -0,0 +1,194 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+from base import Module
+
+class InstanceNorm(Module):
+ r"""Applies instance normalization [1] on the inputs.
+
+ Computes
+
+ .. math::
+
+ y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta,
+
+ where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+ parameters initialized at 1 and 0 respectively. Both are of size :attr:`dims`,
+ if :attr:`affine` is ``True``.
+
+ Args:
+ dims (int): The number of features of the input.
+ eps (float): A value added to the denominator for numerical stability. Default: ``1e-5``.
+ affine (bool): Default: ``False``.
+
+ Shape:
+ - Input: :math:`(..., C)` where :math:`C` is equal to :attr:`dims`.
+ - Output: Same shape as the input.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+ >>> x = mx.random.normal((8, 4, 4, 16))
+ >>> inorm = nn.InstanceNorm(dims=16)
+ >>> output = inorm(x)
+
+ References:
+ [1]: https://arxiv.org/abs/1607.08022
+ """
+ def __init__(self, dims: int, eps: float = ..., affine: bool = ...) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class LayerNorm(Module):
+ r"""Applies layer normalization [1] on the inputs.
+
+ Computes
+
+ .. math::
+
+ y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+ where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+ parameters initialized at 1 and 0 respectively.
+
+ [1]: https://arxiv.org/abs/1607.06450
+
+ Args:
+ dims (int): The feature dimension of the input to normalize over
+ eps (float): A small additive constant for numerical stability
+ affine (bool): If True learn an affine transform to apply after the
+ normalization
+ bias (bool): If True include a translation to the affine
+ transformation. If set to False the transformation is not really affine
+ just scaling.
+ """
+ def __init__(
+ self, dims: int, eps: float = ..., affine: bool = ..., bias: bool = ...
+ ) -> None: ...
+ def __call__(self, x) -> mx.array: ...
+
+class RMSNorm(Module):
+ r"""Applies Root Mean Square normalization [1] to the inputs.
+
+ Computes
+
+ .. math::
+
+ y = \frac{x}{\sqrt{E[x^2] + \epsilon}} \gamma
+
+ where :math:`\gamma` is a learned per feature dimension parameter initialized at
+ 1.
+
+ Note the accumulation for the mean is done in 32-bit precision.
+
+ [1]: https://arxiv.org/abs/1910.07467
+
+ Args:
+ dims (int): The feature dimension of the input to normalize over
+ eps (float): A small additive constant for numerical stability
+ """
+ def __init__(self, dims: int, eps: float = ...) -> None: ...
+ def __call__(self, x) -> mx.array: ...
+
+class GroupNorm(Module):
+ r"""Applies Group Normalization [1] to the inputs.
+
+ Computes the same normalization as layer norm, namely
+
+ .. math::
+
+ y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+ where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+ parameters initialized at 1 and 0 respectively. However, the mean and
+ variance are computed over the spatial dimensions and each group of
+ features. In particular, the input is split into num_groups across the
+ feature dimension.
+
+ The feature dimension is assumed to be the last dimension and the dimensions
+ that precede it (except the first) are considered the spatial dimensions.
+
+ [1]: https://arxiv.org/abs/1803.08494
+
+ Args:
+ num_groups (int): Number of groups to separate the features into
+ dims (int): The feature dimensions of the input to normalize over
+ eps (float): A small additive constant for numerical stability
+ affine (bool): If True learn an affine transform to apply after the
+ normalization.
+ pytorch_compatible (bool): If True perform the group normalization in
+ the same order/grouping as PyTorch.
+ """
+ def __init__(
+ self,
+ num_groups: int,
+ dims: int,
+ eps: float = ...,
+ affine: bool = ...,
+ pytorch_compatible: bool = ...,
+ ) -> None: ...
+ def __call__(self, x) -> mx.array: ...
+
+class BatchNorm(Module):
+ r"""Applies Batch Normalization over a 2D or 3D input.
+
+ Computes
+
+ .. math::
+
+ y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+ where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+ parameters initialized at 1 and 0 respectively.
+
+ The input shape is specified as ``NC`` or ``NLC``, where ``N`` is the
+ batch, ``C`` is the number of features or channels, and ``L`` is the
+ sequence length. The output has the same shape as the input. For
+ four-dimensional arrays, the shape is ``NHWC``, where ``H`` and ``W`` are
+ the height and width respectively.
+
+ For more information on Batch Normalization, see the original paper `Batch
+ Normalization: Accelerating Deep Network Training by Reducing Internal
+ Covariate Shift <https://arxiv.org/abs/1502.03167>`_.
+
+ Args:
+ num_features (int): The feature dimension to normalize over.
+ eps (float, optional): A small additive constant for numerical
+ stability. Default: ``1e-5``.
+ momentum (float, optional): The momentum for updating the running
+ mean and variance. Default: ``0.1``.
+ affine (bool, optional): If ``True``, apply a learned affine
+ transformation after the normalization. Default: ``True``.
+ track_running_stats (bool, optional): If ``True``, track the
+ running mean and variance. Default: ``True``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+ >>> x = mx.random.normal((5, 4))
+ >>> bn = nn.BatchNorm(num_features=4, affine=True)
+ >>> output = bn(x)
+ """
+ def __init__(
+ self,
+ num_features: int,
+ eps: float = ...,
+ momentum: float = ...,
+ affine: bool = ...,
+ track_running_stats: bool = ...,
+ ) -> None: ...
+ def unfreeze(self, *args, **kwargs): # -> None:
+ """Wrap unfreeze to make sure that running_mean and var are always
+ frozen parameters."""
+
+ def __call__(self, x: mx.array) -> mx.array:
+ """
+ Forward pass of BatchNorm.
+
+ Args:
+ x (array): Input tensor.
+
+ Returns:
+ array: Normalized output tensor.
+ """
diff --git a/typings/mlx/nn/layers/pooling.pyi b/typings/mlx/nn/layers/pooling.pyi
new file mode 100644
index 00000000..36b0ca24
--- /dev/null
+++ b/typings/mlx/nn/layers/pooling.pyi
@@ -0,0 +1,242 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Optional, Tuple, Union
+
+import mlx.core as mx
+from base import Module
+
+class _Pool(Module):
+ def __init__(
+ self, pooling_function, kernel_size, stride, padding, padding_value
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class _Pool1d(_Pool):
+ def __init__(
+ self,
+ pooling_function,
+ padding_value,
+ kernel_size: Union[int, Tuple[int]],
+ stride: Optional[Union[int, Tuple[int]]] = ...,
+ padding: Union[int, Tuple[int]] = ...,
+ ) -> None: ...
+
+class _Pool2d(_Pool):
+ def __init__(
+ self,
+ pooling_function,
+ padding_value,
+ kernel_size: Union[int, Tuple[int, int]],
+ stride: Optional[Union[int, Tuple[int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int]]] = ...,
+ ) -> None: ...
+
+class _Pool3d(_Pool):
+ def __init__(
+ self,
+ pooling_function,
+ padding_value,
+ kernel_size: Union[int, Tuple[int, int, int]],
+ stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ ) -> None: ...
+
+class MaxPool1d(_Pool1d):
+ r"""Applies 1-dimensional max pooling.
+
+ Spatially downsamples the input by taking the maximum of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ Args:
+ kernel_size (int or tuple(int)): The size of the pooling window kernel.
+ stride (int or tuple(int), optional): The stride of the pooling window.
+ Default: ``kernel_size``.
+ padding (int or tuple(int), optional): How much negative infinity
+ padding to apply to the input. The padding amount is applied to
+ both sides of the spatial axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(4, 16, 5))
+ >>> pool = nn.MaxPool1d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int]],
+ stride: Optional[Union[int, Tuple[int]]] = ...,
+ padding: Union[int, Tuple[int]] = ...,
+ ) -> None: ...
+
+class AvgPool1d(_Pool1d):
+ r"""Applies 1-dimensional average pooling.
+
+ Spatially downsamples the input by taking the average of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ Args:
+ kernel_size (int or tuple(int)): The size of the pooling window kernel.
+ stride (int or tuple(int), optional): The stride of the pooling window.
+ Default: ``kernel_size``.
+ padding (int or tuple(int), optional): How much zero padding to apply to
+ the input. The padding amount is applied to both sides of the spatial
+ axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(4, 16, 5))
+ >>> pool = nn.AvgPool1d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int]],
+ stride: Optional[Union[int, Tuple[int]]] = ...,
+ padding: Union[int, Tuple[int]] = ...,
+ ) -> None: ...
+
+class MaxPool2d(_Pool2d):
+ r"""Applies 2-dimensional max pooling.
+
+ Spatially downsamples the input by taking the maximum of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
+
+ * a single ``int`` -- in which case the same value is used for both the
+ height and width axis.
+ * a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is
+ used for the height axis, the second ``int`` for the width axis.
+
+ Args:
+ kernel_size (int or tuple(int, int)): The size of the pooling window.
+ stride (int or tuple(int, int), optional): The stride of the pooling
+ window. Default: ``kernel_size``.
+ padding (int or tuple(int, int), optional): How much negative infinity
+ padding to apply to the input. The padding is applied on both sides
+ of the height and width axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(8, 32, 32, 4))
+ >>> pool = nn.MaxPool2d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int, int]],
+ stride: Optional[Union[int, Tuple[int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int]]] = ...,
+ ) -> None: ...
+
+class AvgPool2d(_Pool2d):
+ r"""Applies 2-dimensional average pooling.
+
+ Spatially downsamples the input by taking the average of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
+
+ * a single ``int`` -- in which case the same value is used for both the
+ height and width axis.
+ * a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is
+ used for the height axis, the second ``int`` for the width axis.
+
+ Args:
+ kernel_size (int or tuple(int, int)): The size of the pooling window.
+ stride (int or tuple(int, int), optional): The stride of the pooling
+ window. Default: ``kernel_size``.
+ padding (int or tuple(int, int), optional): How much zero
+ padding to apply to the input. The padding is applied on both sides
+ of the height and width axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(8, 32, 32, 4))
+ >>> pool = nn.AvgPool2d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int, int]],
+ stride: Optional[Union[int, Tuple[int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int]]] = ...,
+ ) -> None: ...
+
+class MaxPool3d(_Pool3d):
+ r"""Applies 3-dimensional max pooling.
+
+ Spatially downsamples the input by taking the maximum of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
+
+ * a single ``int`` -- in which case the same value is used for the depth,
+ height, and width axis.
+ * a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used
+ for the depth axis, the second ``int`` for the height axis, and the third
+ ``int`` for the width axis.
+
+ Args:
+ kernel_size (int or tuple(int, int, int)): The size of the pooling window.
+ stride (int or tuple(int, int, int), optional): The stride of the pooling
+ window. Default: ``kernel_size``.
+ padding (int or tuple(int, int, int), optional): How much negative infinity
+ padding to apply to the input. The padding is applied on both sides
+ of the depth, height and width axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
+ >>> pool = nn.MaxPool3d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int, int, int]],
+ stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ ) -> None: ...
+
+class AvgPool3d(_Pool3d):
+ r"""Applies 3-dimensional average pooling.
+
+ Spatially downsamples the input by taking the average of a sliding window
+ of size ``kernel_size`` and sliding stride ``stride``.
+
+ The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
+
+ * a single ``int`` -- in which case the same value is used for the depth,
+ height, and width axis.
+ * a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used
+ for the depth axis, the second ``int`` for the height axis, and the third
+ ``int`` for the width axis.
+
+ Args:
+ kernel_size (int or tuple(int, int, int)): The size of the pooling window.
+ stride (int or tuple(int, int, int), optional): The stride of the pooling
+ window. Default: ``kernel_size``.
+ padding (int or tuple(int, int, int), optional): How much zero
+ padding to apply to the input. The padding is applied on both sides
+ of the depth, height and width axis. Default: ``0``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import layers as nn
+ >>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
+ >>> pool = nn.AvgPool3d(kernel_size=2, stride=2)
+ >>> pool(x)
+ """
+ def __init__(
+ self,
+ kernel_size: Union[int, Tuple[int, int, int]],
+ stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
+ ) -> None: ...
diff --git a/typings/mlx/nn/layers/positional_encoding.pyi b/typings/mlx/nn/layers/positional_encoding.pyi
new file mode 100644
index 00000000..14e07e14
--- /dev/null
+++ b/typings/mlx/nn/layers/positional_encoding.pyi
@@ -0,0 +1,80 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Optional
+
+import mlx.core as mx
+from base import Module
+
+class RoPE(Module):
+ """Implements the rotary positional encoding.
+
+ The traditional implementation rotates consecutive pairs of elements in the
+ feature dimension while the default implementation rotates pairs with
+ stride half the feature dimensions for efficiency.
+
+ For more details see `RoFormer: Enhanced Transformer with Rotary Position
+ Embedding <https://arxiv.org/abs/2104.09864>`_.
+
+ Args:
+ dims (int): The feature dimensions to be rotated. If the input feature
+ is larger than dims then the rest is left unchanged.
+ traditional (bool, optional): If set to ``True`` choose the traditional
+ implementation which is slightly less efficient. Default: ``False``.
+ base (float, optional): The base used to compute angular frequency for
+ each dimension in the positional encodings. Default: ``10000``.
+ scale (float, optional): The scale used to scale the positions. Default: ``1.0``.
+ """
+ def __init__(
+ self, dims: int, traditional: bool = ..., base: float = ..., scale: float = ...
+ ) -> None: ...
+ def __call__(self, x, offset: int = ...) -> mx.array: ...
+
+class SinusoidalPositionalEncoding(Module):
+ r"""Implements sinusoidal positional encoding.
+
+ For more details see the paper `Attention Is All You Need
+ <https://arxiv.org/abs/1706.03762>`_.
+
+ Args:
+ dims (int): The dimensionality of the resulting positional embeddings.
+ min_freq (float, optional): The minimum frequency expected. Default:
+ ``0.0001``.
+ max_freq (float, optional): The maximum frequency expected. Default:
+ ``1``.
+ scale (float, optional): A multiplicative scale for the embeddings.
+ Default: ``sqrt(2/dims)``.
+ cos_first (bool, optional): If ``True`` embed using ``[cos(x); sin(x)]``
+ instead of the reverse. Default: ``False``.
+ full_turns (bool, optional): If ``True`` multiply the frequencies with
+ :math:`2\pi`. Default: ``False``.
+ """
+ def __init__(
+ self,
+ dims: int,
+ min_freq: float = ...,
+ max_freq: float = ...,
+ scale: Optional[float] = ...,
+ cos_first: bool = ...,
+ full_turns: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class ALiBi(Module):
+ _alibi_mask_key = ...
+ _alibi_mask = ...
+ @classmethod
+ def create_alibi_matrix(
+ cls,
+ q_sequence_length: int,
+ k_sequence_length: int,
+ num_heads: int,
+ offset: int,
+ dtype=...,
+ ) -> mx.array | None: ...
+ @staticmethod
+ def create_alibi_slope(num_heads: int) -> mx.array: ...
+ def __call__(
+ self, attention_scores: mx.array, offset=..., mask=...
+ ) -> mx.array: ...
diff --git a/typings/mlx/nn/layers/quantized.pyi b/typings/mlx/nn/layers/quantized.pyi
new file mode 100644
index 00000000..137a4c8e
--- /dev/null
+++ b/typings/mlx/nn/layers/quantized.pyi
@@ -0,0 +1,125 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Callable, Optional, Union
+
+import mlx.core as mx
+from base import Module
+
+def quantize(
+ model: Module,
+ group_size: int = ...,
+ bits: int = ...,
+ *,
+ mode: str = ...,
+ class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
+): # -> None:
+ """Quantize the sub-modules of a module according to a predicate.
+
+ By default all layers that define a ``to_quantized(group_size, bits)``
+ method will be quantized. Both :obj:`Linear` and :obj:`Embedding` layers
+ will be quantized. Note also, the module is updated in-place.
+
+ Args:
+ model (Module): The model whose leaf modules may be quantized.
+ group_size (int): The quantization group size (see
+ :func:`mlx.core.quantize`). Default: ``64``.
+ bits (int): The number of bits per parameter (see
+ :func:`mlx.core.quantize`). Default: ``4``.
+ mode (str): The quantization method to use (see
+ :func:`mlx.core.quantize`). Default: ``"affine"``.
+ class_predicate (Optional[Callable]): A callable which receives the
+ :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a
+ dict of params for `to_quantized` if it should be quantized and
+ ``False`` otherwise. If ``None``, then all layers that define a
+ ``to_quantized(group_size, bits)`` method are quantized.
+ Default: ``None``.
+ """
+
+class QuantizedEmbedding(Module):
+ """The same as :obj:`Embedding` but with a quantized weight matrix.
+
+ :obj:`QuantizedEmbedding` also provides a :meth:`from_embedding`
+ classmethod to convert embedding layers to :obj:`QuantizedEmbedding`
+ layers.
+
+ Args:
+ num_embeddings (int): How many possible discrete tokens can we embed.
+ Usually called the vocabulary size.
+ dims (int): The dimensionality of the embeddings.
+ group_size (int, optional): The group size to use for the quantized
+ weight. See :func:`~mlx.core.quantize`. Default: ``64``.
+ bits (int, optional): The bit width to use for the quantized weight.
+ See :func:`~mlx.core.quantize`. Default: ``4``.
+ mode (str): The quantization method to use (see
+ :func:`mlx.core.quantize`). Default: ``"affine"``.
+ """
+ def __init__(
+ self,
+ num_embeddings: int,
+ dims: int,
+ group_size: int = ...,
+ bits: int = ...,
+ mode: str = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ def as_linear(self, x: mx.array) -> mx.array:
+ """
+ Call the quantized embedding layer as a quantized linear layer.
+
+ Use this for example when input embedding and output projection
+ weights are tied.
+ """
+
+ @classmethod
+ def from_embedding(
+ cls,
+ embedding_layer: Module,
+ group_size: int = ...,
+ bits: int = ...,
+ mode: str = ...,
+ ) -> QuantizedEmbedding:
+ """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
+
+class QuantizedLinear(Module):
+ """Applies an affine transformation to the input using a quantized weight matrix.
+
+ It is the quantized equivalent of :class:`Linear`. For now its
+ parameters are frozen and will not be included in any gradient computation
+ but this will probably change in the future.
+
+ :obj:`QuantizedLinear` also provides a classmethod :meth:`from_linear` to
+ convert linear layers to :obj:`QuantizedLinear` layers.
+
+ Args:
+ input_dims (int): The dimensionality of the input features.
+ output_dims (int): The dimensionality of the output features.
+ bias (bool, optional): If set to ``False`` then the layer will not use
+ a bias. Default: ``True``.
+ group_size (int, optional): The group size to use for the quantized
+ weight. See :func:`~mlx.core.quantize`. Default: ``64``.
+ bits (int, optional): The bit width to use for the quantized weight.
+ See :func:`~mlx.core.quantize`. Default: ``4``.
+ mode (str): The quantization method to use (see
+ :func:`mlx.core.quantize`). Default: ``"affine"``.
+ """
+ def __init__(
+ self,
+ input_dims: int,
+ output_dims: int,
+ bias: bool = ...,
+ group_size: int = ...,
+ bits: int = ...,
+ mode: str = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+ @classmethod
+ def from_linear(
+ cls,
+ linear_layer: Module,
+ group_size: int = ...,
+ bits: int = ...,
+ mode: str = ...,
+ ) -> QuantizedLinear:
+ """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""
diff --git a/typings/mlx/nn/layers/recurrent.pyi b/typings/mlx/nn/layers/recurrent.pyi
new file mode 100644
index 00000000..d31d9382
--- /dev/null
+++ b/typings/mlx/nn/layers/recurrent.pyi
@@ -0,0 +1,113 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Callable, Optional
+
+import mlx.core as mx
+from base import Module
+
+class RNN(Module):
+ r"""An Elman recurrent layer.
+
+ The input is a sequence of shape ``NLD`` or ``LD`` where:
+
+ * ``N`` is the optional batch dimension
+ * ``L`` is the sequence length
+ * ``D`` is the input's feature dimension
+
+ Concretely, for each element along the sequence length axis, this
+ layer applies the function:
+
+ .. math::
+
+ h_{t + 1} = \text{tanh} (W_{ih}x_t + W_{hh}h_t + b)
+
+ The hidden state :math:`h` has shape ``NH`` or ``H``, depending on
+ whether the input is batched or not. Returns the hidden state at each
+ time step, of shape ``NLH`` or ``LH``.
+
+ Args:
+ input_size (int): Dimension of the input, ``D``.
+ hidden_size (int): Dimension of the hidden state, ``H``.
+ bias (bool, optional): Whether to use a bias. Default: ``True``.
+ nonlinearity (callable, optional): Non-linearity to use. If ``None``,
+ then func:`tanh` is used. Default: ``None``.
+ """
+ def __init__(
+ self,
+ input_size: int,
+ hidden_size: int,
+ bias: bool = ...,
+ nonlinearity: Optional[Callable] = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array, hidden=...) -> mx.array: ...
+
+class GRU(Module):
+ r"""A gated recurrent unit (GRU) RNN layer.
+
+ The input has shape ``NLD`` or ``LD`` where:
+
+ * ``N`` is the optional batch dimension
+ * ``L`` is the sequence length
+ * ``D`` is the input's feature dimension
+
+ Concretely, for each element of the sequence, this layer computes:
+
+ .. math::
+
+ \begin{aligned}
+ r_t &= \sigma (W_{xr}x_t + W_{hr}h_t + b_{r}) \\
+ z_t &= \sigma (W_{xz}x_t + W_{hz}h_t + b_{z}) \\
+ n_t &= \text{tanh}(W_{xn}x_t + b_{n} + r_t \odot (W_{hn}h_t + b_{hn})) \\
+ h_{t + 1} &= (1 - z_t) \odot n_t + z_t \odot h_t
+ \end{aligned}
+
+ The hidden state :math:`h` has shape ``NH`` or ``H`` depending on
+ whether the input is batched or not. Returns the hidden state at each
+ time step of shape ``NLH`` or ``LH``.
+
+ Args:
+ input_size (int): Dimension of the input, ``D``.
+ hidden_size (int): Dimension of the hidden state, ``H``.
+ bias (bool): Whether to use biases or not. Default: ``True``.
+ """
+ def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ...
+ def __call__(self, x: mx.array, hidden=...) -> mx.array: ...
+
+class LSTM(Module):
+ r"""An LSTM recurrent layer.
+
+ The input has shape ``NLD`` or ``LD`` where:
+
+ * ``N`` is the optional batch dimension
+ * ``L`` is the sequence length
+ * ``D`` is the input's feature dimension
+
+ Concretely, for each element of the sequence, this layer computes:
+
+ .. math::
+ \begin{aligned}
+ i_t &= \sigma (W_{xi}x_t + W_{hi}h_t + b_{i}) \\
+ f_t &= \sigma (W_{xf}x_t + W_{hf}h_t + b_{f}) \\
+ g_t &= \text{tanh} (W_{xg}x_t + W_{hg}h_t + b_{g}) \\
+ o_t &= \sigma (W_{xo}x_t + W_{ho}h_t + b_{o}) \\
+ c_{t + 1} &= f_t \odot c_t + i_t \odot g_t \\
+ h_{t + 1} &= o_t \text{tanh}(c_{t + 1})
+ \end{aligned}
+
+ The hidden state :math:`h` and cell state :math:`c` have shape ``NH``
+ or ``H``, depending on whether the input is batched or not.
+
+ The layer returns two arrays, the hidden state and the cell state at
+ each time step, both of shape ``NLH`` or ``LH``.
+
+ Args:
+ input_size (int): Dimension of the input, ``D``.
+ hidden_size (int): Dimension of the hidden state, ``H``.
+ bias (bool): Whether to use biases or not. Default: ``True``.
+ """
+ def __init__(self, input_size: int, hidden_size: int, bias: bool = ...) -> None: ...
+ def __call__(
+ self, x: mx.array, hidden=..., cell=...
+ ) -> tuple[mx.array, mx.array]: ...
diff --git a/typings/mlx/nn/layers/transformer.pyi b/typings/mlx/nn/layers/transformer.pyi
new file mode 100644
index 00000000..9274a823
--- /dev/null
+++ b/typings/mlx/nn/layers/transformer.pyi
@@ -0,0 +1,168 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Any, Callable, Optional
+
+import mlx.core as mx
+from base import Module
+
+class MultiHeadAttention(Module):
+ """Implements the scaled dot product attention with multiple heads.
+
+ Given inputs for queries, keys and values the ``MultiHeadAttention``
+ produces new values by aggregating information from the input values
+ according to the similarities of the input queries and keys.
+
+ All inputs as well as the output are linearly projected without biases by
+ default.
+
+ ``MultiHeadAttention`` also takes an optional additive attention mask that
+ should be broadcastable with ``(batch, num_heads, # queries, # keys)``. The
+ mask should have ``-inf`` or very large negative numbers at the positions
+ that should *not* be attended to.
+
+ Args:
+ dims (int): The model dimensions. This is also the default
+ value for the queries, keys, values, and the output.
+ num_heads (int): The number of attention heads to use.
+ query_input_dims (int, optional): The input dimensions of the queries.
+ Default: ``dims``.
+ key_input_dims (int, optional): The input dimensions of the keys.
+ Default: ``dims``.
+ value_input_dims (int, optional): The input dimensions of the values.
+ Default: ``key_input_dims``.
+ value_dims (int, optional): The dimensions of the values after the
+ projection. Default: ``dims``.
+ value_output_dims (int, optional): The dimensions the new values will
+ be projected to. Default: ``dims``.
+ bias (bool, optional): Whether or not to use a bias in the projections.
+ Default: ``False``.
+ """
+ def __init__(
+ self,
+ dims: int,
+ num_heads: int,
+ query_input_dims: Optional[int] = ...,
+ key_input_dims: Optional[int] = ...,
+ value_input_dims: Optional[int] = ...,
+ value_dims: Optional[int] = ...,
+ value_output_dims: Optional[int] = ...,
+ bias: bool = ...,
+ ) -> None: ...
+ def __call__(
+ self, queries: mx.array, keys: mx.array, values: mx.array, mask: mx.array = ...
+ ) -> mx.array: ...
+ @staticmethod
+ def create_additive_causal_mask(N: int, dtype: mx.Dtype = ...) -> mx.array: ...
+
+class TransformerEncoderLayer(Module):
+ def __init__(
+ self,
+ dims: int,
+ num_heads: int,
+ mlp_dims: Optional[int] = ...,
+ dropout: float = ...,
+ activation: Callable[[Any], Any] = ...,
+ norm_first: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ...
+
+class TransformerEncoder(Module):
+ def __init__(
+ self,
+ num_layers: int,
+ dims: int,
+ num_heads: int,
+ mlp_dims: Optional[int] = ...,
+ dropout: float = ...,
+ activation=...,
+ norm_first: bool = ...,
+ checkpoint: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array, mask: mx.array) -> mx.array: ...
+
+class TransformerDecoderLayer(Module):
+ def __init__(
+ self,
+ dims: int,
+ num_heads: int,
+ mlp_dims: Optional[int] = ...,
+ dropout: float = ...,
+ activation: Callable[[Any], Any] = ...,
+ norm_first: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ...
+
+class TransformerDecoder(Module):
+ def __init__(
+ self,
+ num_layers: int,
+ dims: int,
+ num_heads: int,
+ mlp_dims: Optional[int] = ...,
+ dropout: float = ...,
+ activation=...,
+ norm_first: bool = ...,
+ checkpoint: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array, memory, x_mask, memory_mask) -> mx.array: ...
+
+class Transformer(Module):
+ """
+ Implements a standard Transformer model.
+
+ The implementation is based on `Attention Is All You Need
+ <https://arxiv.org/abs/1706.03762>`_.
+
+ The Transformer model contains an encoder and a decoder. The encoder
+ processes the input sequence and the decoder generates the output sequence.
+ The interaction between encoder and decoder happens through the attention
+ mechanism.
+
+ Args:
+ dims (int, optional): The number of expected features in the
+ encoder/decoder inputs. Default: ``512``.
+ num_heads (int, optional): The number of attention heads. Default:
+ ``8``.
+ num_encoder_layers (int, optional): The number of encoder layers in the
+ Transformer encoder. Default: ``6``.
+ num_decoder_layers (int, optional): The number of decoder layers in the
+ Transformer decoder. Default: ``6``.
+ mlp_dims (int, optional): The hidden dimension of the MLP block in each
+ Transformer layer. Defaults to ``4*dims`` if not provided. Default:
+ ``None``.
+ dropout (float, optional): The dropout value for the Transformer
+ encoder and decoder. Dropout is used after each attention layer and
+ the activation in the MLP layer. Default: ``0.0``.
+ activation (function, optional): the activation function for the MLP
+ hidden layer. Default: :func:`relu`.
+ custom_encoder (nn.Module, optional): A custom encoder to replace the
+ standard Transformer encoder. Default: ``None``.
+ custom_decoder (nn.Module, optional): A custom decoder to replace the
+ standard Transformer decoder. Default: ``None``.
+ norm_first (bool, optional): if ``True``, encoder and decoder layers
+ will perform layer normalization before attention and MLP
+ operations, otherwise after. Default: ``True``.
+ checkpoint (bool, optional): if ``True`` perform gradient checkpointing
+ to reduce the memory usage at the expense of more computation.
+ Default: ``False``.
+ """
+ def __init__(
+ self,
+ dims: int = ...,
+ num_heads: int = ...,
+ num_encoder_layers: int = ...,
+ num_decoder_layers: int = ...,
+ mlp_dims: Optional[int] = ...,
+ dropout: float = ...,
+ activation: Callable[[Any], Any] = ...,
+ custom_encoder: Optional[Any] = ...,
+ custom_decoder: Optional[Any] = ...,
+ norm_first: bool = ...,
+ checkpoint: bool = ...,
+ ) -> None: ...
+ def __call__(
+ self, src, tgt, src_mask, tgt_mask, memory_mask
+ ) -> mx.array: # -> array | Any:
+ ...
diff --git a/typings/mlx/nn/layers/upsample.pyi b/typings/mlx/nn/layers/upsample.pyi
new file mode 100644
index 00000000..1ef3298c
--- /dev/null
+++ b/typings/mlx/nn/layers/upsample.pyi
@@ -0,0 +1,87 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Literal, Tuple, Union
+
+import mlx.core as mx
+from base import Module
+
+def upsample_nearest(x: mx.array, scale_factor: Tuple) -> mx.array: ...
+def upsample_linear(
+ x: mx.array, scale_factor: Tuple, align_corners: bool = ...
+): # -> int:
+ ...
+def upsample_cubic(
+ x: mx.array, scale_factor: Tuple, align_corners: bool = ...
+): # -> int:
+ ...
+
+class Upsample(Module):
+ r"""Upsample the input signal spatially.
+
+ The spatial dimensions are by convention dimensions ``1`` to ``x.ndim -
+ 2``. The first is the batch dimension and the last is the feature
+ dimension.
+
+ For example, an audio signal would be 3D with 1 spatial dimension, an image
+ 4D with 2 and so on and so forth.
+
+ There are three upsampling algorithms implemented nearest neighbor upsampling,
+ linear interpolation, and cubic interpolation. All can be applied to any number
+ of spatial dimensions. The linear interpolation will be bilinear, trilinear etc
+ when applied to more than one spatial dimension. And cubic interpolation will be
+ bicubic when there are 2 spatial dimensions.
+
+ .. note::
+ When using one of the linear or cubic interpolation modes the ``align_corners``
+ argument changes how the corners are treated in the input image. If
+ ``align_corners=True`` then the top and left edge of the input and
+ output will be matching as will the bottom right edge.
+
+ Parameters:
+ scale_factor (float or tuple): The multiplier for the spatial size.
+ If a ``float`` is provided, it is the multiplier for all spatial dimensions.
+ Otherwise, the number of scale factors provided must match the
+ number of spatial dimensions.
+ mode (str, optional): The upsampling algorithm, either ``"nearest"``,
+ ``"linear"`` or ``"cubic"``. Default: ``"nearest"``.
+ align_corners (bool, optional): Changes the way the corners are treated
+ during ``"linear"`` and ``"cubic"`` upsampling. See the note above and the
+ examples below for more details. Default: ``False``.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+ >>> x = mx.arange(1, 5).reshape((1, 2, 2, 1))
+ >>> x
+ array([[[[1],
+ [2]],
+ [[3],
+ [4]]]], dtype=int32)
+ >>> n = nn.Upsample(scale_factor=2, mode='nearest')
+ >>> n(x).squeeze()
+ array([[1, 1, 2, 2],
+ [1, 1, 2, 2],
+ [3, 3, 4, 4],
+ [3, 3, 4, 4]], dtype=int32)
+ >>> b = nn.Upsample(scale_factor=2, mode='linear')
+ >>> b(x).squeeze()
+ array([[1, 1.25, 1.75, 2],
+ [1.5, 1.75, 2.25, 2.5],
+ [2.5, 2.75, 3.25, 3.5],
+ [3, 3.25, 3.75, 4]], dtype=float32)
+ >>> b = nn.Upsample(scale_factor=2, mode='linear', align_corners=True)
+ >>> b(x).squeeze()
+ array([[1, 1.33333, 1.66667, 2],
+ [1.66667, 2, 2.33333, 2.66667],
+ [2.33333, 2.66667, 3, 3.33333],
+ [3, 3.33333, 3.66667, 4]], dtype=float32)
+ """
+ def __init__(
+ self,
+ scale_factor: Union[float, Tuple],
+ mode: Literal["nearest", "linear", "cubic"] = ...,
+ align_corners: bool = ...,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
diff --git a/typings/mlx/nn/losses.pyi b/typings/mlx/nn/losses.pyi
new file mode 100644
index 00000000..9b5ded9e
--- /dev/null
+++ b/typings/mlx/nn/losses.pyi
@@ -0,0 +1,419 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Literal, Optional
+
+import mlx.core as mx
+
+Reduction = Literal["none", "mean", "sum"]
+
+def cross_entropy(
+ logits: mx.array,
+ targets: mx.array,
+ weights: Optional[mx.array] = ...,
+ axis: int = ...,
+ label_smoothing: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ """
+ Computes the cross entropy loss.
+
+ Args:
+ logits (array): The unnormalized logits.
+ targets (array): The ground truth values. These can be class indices or
+ probabilities for each class. If the ``targets`` are class indices,
+ then ``targets`` shape should match the ``logits`` shape with
+ the ``axis`` dimension removed. If the ``targets`` are probabilities
+ (or one-hot encoded), then the ``targets`` shape should be the same as
+ the ``logits`` shape.
+ weights (array, optional): Optional weights for each target. Default: ``None``.
+ axis (int, optional): The axis over which to compute softmax. Default: ``-1``.
+ label_smoothing (float, optional): Label smoothing factor. Default: ``0``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed cross entropy loss.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+ >>>
+ >>> # Class indices as targets
+ >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+ >>> targets = mx.array([0, 1])
+ >>> nn.losses.cross_entropy(logits, targets)
+ array([0.0485873, 0.0485873], dtype=float32)
+ >>>
+ >>> # Probabilities (or one-hot vectors) as targets
+ >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]])
+ >>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]])
+ >>> nn.losses.cross_entropy(logits, targets)
+ array([0.348587, 0.348587], dtype=float32)
+ """
+
+def binary_cross_entropy(
+ inputs: mx.array,
+ targets: mx.array,
+ weights: Optional[mx.array] = ...,
+ with_logits: bool = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ """
+ Computes the binary cross entropy loss.
+
+ By default, this function takes the pre-sigmoid logits, which results in a faster
+ and more precise loss. For improved numerical stability when ``with_logits=False``,
+ the loss calculation clips the input probabilities (in log-space) to a minimum value
+ of ``-100``.
+
+ Args:
+ inputs (array): The predicted values. If ``with_logits`` is ``True``, then
+ ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities.
+ targets (array): The binary target values in {0, 1}.
+ with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``.
+ weights (array, optional): Optional weights for each target. Default: ``None``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+ Returns:
+ array: The computed binary cross entropy loss.
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+
+ >>> logits = mx.array([0.105361, 0.223144, 1.20397, 0.916291])
+ >>> targets = mx.array([0, 0, 1, 1])
+ >>> loss = nn.losses.binary_cross_entropy(logits, targets, reduction="mean")
+ >>> loss
+ array(0.539245, dtype=float32)
+
+ >>> probs = mx.array([0.1, 0.1, 0.4, 0.4])
+ >>> targets = mx.array([0, 0, 1, 1])
+ >>> loss = nn.losses.binary_cross_entropy(probs, targets, with_logits=False, reduction="mean")
+ >>> loss
+ array(0.510826, dtype=float32)
+ """
+
+def l1_loss(
+ predictions: mx.array, targets: mx.array, reduction: Reduction = ...
+) -> mx.array:
+ """
+ Computes the L1 loss.
+
+ Args:
+ predictions (array): The predicted values.
+ targets (array): The target values.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+ Returns:
+ array: The computed L1 loss.
+ """
+
+def mse_loss(
+ predictions: mx.array, targets: mx.array, reduction: Reduction = ...
+) -> mx.array:
+ """
+ Computes the mean squared error loss.
+
+ Args:
+ predictions (array): The predicted values.
+ targets (array): The target values.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+ Returns:
+ array: The computed mean squared error loss.
+ """
+
+def nll_loss(
+ inputs: mx.array, targets: mx.array, axis: int = ..., reduction: Reduction = ...
+) -> mx.array:
+ """
+ Computes the negative log likelihood loss.
+
+ Args:
+ inputs (array): The predicted distribution in log space.
+ targets (array): The target values.
+ axis (int, optional): The distribution axis. Default: ``-1``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed NLL loss.
+ """
+
+def gaussian_nll_loss(
+ inputs: mx.array,
+ targets: mx.array,
+ vars: mx.array,
+ full: bool = ...,
+ eps: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ r"""
+ Computes the negative log likelihood loss for a Gaussian distribution.
+
+ The loss is given by:
+
+ .. math::
+ \frac{1}{2}\left(\log\left(\max\left(\text{vars},
+ \ \epsilon\right)\right) + \frac{\left(\text{inputs} - \text{targets} \right)^2}
+ {\max\left(\text{vars}, \ \epsilon \right)}\right) + \text{const.}
+
+ where ``inputs`` are the predicted means and ``vars`` are the the
+ predicted variances.
+
+ Args:
+ inputs (array): The predicted expectation of the Gaussian distribution.
+ targets (array): The target values (samples from the Gaussian distribution).
+ vars (array): The predicted variance of the Gaussian distribution.
+ full (bool, optional): Whether to include the constant term in the loss calculation.
+ Default: ``False``.
+ eps (float, optional): Small positive constant for numerical stability.
+ Default: ``1e-6``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The Gaussian NLL loss.
+ """
+
+def kl_div_loss(
+ inputs: mx.array, targets: mx.array, axis: int = ..., reduction: Reduction = ...
+) -> mx.array:
+ """
+ Computes the Kullback-Leibler divergence loss.
+
+ Computes the following when ``reduction == 'none'``:
+
+ .. code-block:: python
+
+ mx.exp(targets) * (targets - inputs).sum(axis)
+
+ Args:
+ inputs (array): Log probabilities for the predicted distribution.
+ targets (array): Log probabilities for the target distribution.
+ axis (int, optional): The distribution axis. Default: ``-1``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed Kullback-Leibler divergence loss.
+ """
+
+def smooth_l1_loss(
+ predictions: mx.array,
+ targets: mx.array,
+ beta: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ r"""
+ Computes the smooth L1 loss.
+
+ The smooth L1 loss is a variant of the L1 loss which replaces the absolute
+ difference with a squared difference when the absolute difference is less
+ than ``beta``.
+
+ The formula for the smooth L1 Loss is:
+
+ .. math::
+
+ l = \begin{cases}
+ 0.5 (x - y)^2 / \beta, & \text{if } |x - y| < \beta \\
+ |x - y| - 0.5 \beta, & \text{otherwise}
+ \end{cases}
+
+ Args:
+ predictions (array): Predicted values.
+ targets (array): Ground truth values.
+ beta (float, optional): The threshold after which the loss changes
+ from the squared to the absolute difference. Default: ``1.0``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+ Returns:
+ array: The computed smooth L1 loss.
+ """
+
+def triplet_loss(
+ anchors: mx.array,
+ positives: mx.array,
+ negatives: mx.array,
+ axis: int = ...,
+ p: int = ...,
+ margin: float = ...,
+ eps: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ r"""
+ Computes the triplet loss for a set of anchor, positive, and negative samples.
+ Margin is represented with alpha in the math section.
+
+ .. math::
+
+ \max\left(\|A - P\|_p - \|A - N\|_p + \alpha, 0\right)
+
+ Args:
+ anchors (array): The anchor samples.
+ positives (array): The positive samples.
+ negatives (array): The negative samples.
+ axis (int, optional): The distribution axis. Default: ``-1``.
+ p (int, optional): The norm degree for pairwise distance. Default: ``2``.
+ margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``.
+ eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input;
+ if reduction is "mean" or "sum", returns a scalar tensor.
+ """
+
+def hinge_loss(
+ inputs: mx.array, targets: mx.array, reduction: Reduction = ...
+) -> mx.array:
+ r"""
+ Computes the hinge loss between inputs and targets.
+
+ .. math::
+
+ \text{hinge}(y, y_{\text{pred}}) = \max(0, 1 - y \cdot y_{\text{pred}})
+
+
+ Args:
+ inputs (array): The predicted values.
+ targets (array): The target values. They should be -1 or 1.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed hinge loss.
+ """
+
+def huber_loss(
+ inputs: mx.array, targets: mx.array, delta: float = ..., reduction: Reduction = ...
+) -> mx.array:
+ r"""
+ Computes the Huber loss between inputs and targets.
+
+ .. math::
+
+ l_{\delta}(a) =
+ \left\{ \begin{array}{ll}
+ \frac{1}{2} a^2 & \text{for } |a| \leq \delta, \\
+ \delta \left( |a| - \frac{1}{2} \delta \right) & \text{otherwise.}
+ \end{array} \right.
+
+ Args:
+ inputs (array): The predicted values.
+ targets (array): The target values.
+ delta (float, optional): The threshold at which to change between L1 and L2 loss.
+ Default: ``1.0``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed Huber loss.
+ """
+
+def log_cosh_loss(
+ inputs: mx.array, targets: mx.array, reduction: Reduction = ...
+) -> mx.array:
+ r"""
+ Computes the log cosh loss between inputs and targets.
+
+ Logcosh acts like L2 loss for small errors, ensuring stable gradients,
+ and like the L1 loss for large errors, reducing sensitivity to outliers. This
+ dual behavior offers a balanced, robust approach for regression tasks.
+
+ .. math::
+
+ \text{logcosh}(y_{\text{true}}, y_{\text{pred}}) =
+ \frac{1}{n} \sum_{i=1}^{n}
+ \log(\cosh(y_{\text{pred}}^{(i)} - y_{\text{true}}^{(i)}))
+
+
+ Args:
+ inputs (array): The predicted values.
+ targets (array): The target values.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed log cosh loss.
+ """
+
+def cosine_similarity_loss(
+ x1: mx.array,
+ x2: mx.array,
+ axis: int = ...,
+ eps: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ r"""
+ Computes the cosine similarity between the two inputs.
+
+ The cosine similarity loss is given by
+
+ .. math::
+
+ \frac{x_1 \cdot x_2}{\max(\|x_1\| \cdot \|x_2\|, \epsilon)}
+
+ Args:
+ x1 (mx.array): The first set of inputs.
+ x2 (mx.array): The second set of inputs.
+ axis (int, optional): The embedding axis. Default: ``1``.
+ eps (float, optional): The minimum value of the denominator used for
+ numerical stability. Default: ``1e-8``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ mx.array: The computed cosine similarity loss.
+ """
+
+def margin_ranking_loss(
+ inputs1: mx.array,
+ inputs2: mx.array,
+ targets: mx.array,
+ margin: float = ...,
+ reduction: Reduction = ...,
+) -> mx.array:
+ r"""
+ Calculate the margin ranking loss that loss given inputs :math:`x_1`, :math:`x_2` and a label
+ :math:`y` (containing 1 or -1).
+
+ The loss is given by:
+
+ .. math::
+ \text{loss} = \max (0, -y * (x_1 - x_2) + \text{margin})
+
+ Where :math:`y` represents ``targets``, :math:`x_1` represents ``inputs1`` and :math:`x_2`
+ represents ``inputs2``.
+
+ Args:
+ inputs1 (array): Scores for the first input.
+ inputs2 (array): Scores for the second input.
+ targets (array): Labels indicating whether samples in ``inputs1`` should be ranked higher
+ than samples in ``inputs2``. Values should be 1 or -1.
+ margin (float, optional): The margin by which the scores should be separated.
+ Default: ``0.0``.
+ reduction (str, optional): Specifies the reduction to apply to the output:
+ ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+ Returns:
+ array: The computed margin ranking loss.
+
+ Examples:
+ >>> import mlx.core as mx
+ >>> import mlx.nn as nn
+ >>> targets = mx.array([1, 1, -1])
+ >>> inputs1 = mx.array([-0.573409, -0.765166, -0.0638])
+ >>> inputs2 = mx.array([0.75596, 0.225763, 0.256995])
+ >>> loss = nn.losses.margin_ranking_loss(inputs1, inputs2, targets)
+ >>> loss
+ array(0.773433, dtype=float32)
+ """
diff --git a/typings/mlx/nn/utils.pyi b/typings/mlx/nn/utils.pyi
new file mode 100644
index 00000000..7df93f12
--- /dev/null
+++ b/typings/mlx/nn/utils.pyi
@@ -0,0 +1,73 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Any, Callable, Optional
+
+import mlx.core as mx
+
+from .layers.base import Module
+
+def value_and_grad(
+ model: Module, fn: Callable
+): # -> _Wrapped[..., Any, ..., tuple[Any, Any]]:
+ """Transform the passed function ``fn`` to a function that computes the
+ gradients of ``fn`` wrt the model's trainable parameters and also its
+ value.
+
+ Args:
+ model (Module): The model whose trainable parameters to compute
+ gradients for
+ fn (Callable): The scalar function to compute gradients for
+
+ Returns:
+ A callable that returns the value of ``fn`` and the gradients wrt the
+ trainable parameters of ``model``
+ """
+
+def checkpoint(
+ module: Module, fn: Optional[Callable] = ...
+): # -> _Wrapped[..., Any, ..., Any]:
+ """Transform the passed callable to one that performs gradient
+ checkpointing with respect to the trainable parameters of the module (and
+ the callable's inputs).
+
+ Args:
+ module (Module): The module for whose parameters we will be
+ performing gradient checkpointing.
+ fn (Callable, optional): The function to checkpoint. If not provided it
+ defaults to the provided module.
+
+ Returns:
+ A callable that saves the inputs and outputs during the forward pass
+ and recomputes all intermediate states during the backward pass.
+ """
+
+def average_gradients(
+ gradients: Any,
+ group: Optional[mx.distributed.Group] = ...,
+ all_reduce_size: int = ...,
+ communication_type: Optional[mx.Dtype] = ...,
+ communication_stream: Optional[mx.Stream] = ...,
+): # -> Any:
+ """Average the gradients across the distributed processes in the passed group.
+
+ This helper enables concatenating several gradients of small arrays to one
+ big all reduce call for better networking performance.
+
+ Args:
+ gradients (Any): The Python tree containing the gradients (it should
+ have the same structure across processes)
+ group (Optional[mlx.core.distributed.Group]): The group of processes to
+ average the gradients. If set to ``None`` the global group is used.
+ Default: ``None``.
+ all_reduce_size (int): Group arrays until their size in bytes exceeds
+ this number. Perform one communication step per group of arrays. If
+ less or equal to 0 array grouping is disabled. Default: ``32MiB``.
+ communication_type (Optional[mlx.core.Dtype]): If provided cast to this
+ type before performing the communication. Typically cast to a
+ smaller float to reduce the communication size. Default: ``None``.
+ communication_stream (Optional[mlx.core.Stream]): The stream to usse
+ for the communication. If unspecified the default communication
+ stream is used which can vary by back-end. Default: ``None``.
+ """
diff --git a/typings/mlx/utils.pyi b/typings/mlx/utils.pyi
new file mode 100644
index 00000000..d005e8cd
--- /dev/null
+++ b/typings/mlx/utils.pyi
@@ -0,0 +1,182 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+def tree_map(
+ 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.
+
+ If ``rest`` is provided, every item is assumed to be a superset of ``tree``
+ and the corresponding leaves are provided as extra positional arguments to
+ ``fn``. In that respect, :meth:`tree_map` is closer to :func:`itertools.starmap`
+ than to :func:`map`.
+
+ The keyword argument ``is_leaf`` decides what constitutes a leaf from
+ ``tree`` similar to :func:`tree_flatten`.
+
+ .. code-block:: python
+
+ import mlx.nn as nn
+ from mlx.utils import tree_map
+
+ model = nn.Linear(10, 10)
+ print(model.parameters().keys())
+ # dict_keys(['weight', 'bias'])
+
+ # square the parameters
+ model.update(tree_map(lambda x: x*x, model.parameters()))
+
+ Args:
+ fn (callable): The function that processes the leaves of the tree.
+ tree (Any): The main Python tree that will be iterated upon.
+ rest (tuple[Any]): Extra trees to be iterated together with ``tree``.
+ is_leaf (callable, optional): An optional callable that returns ``True``
+ if the passed object is considered a leaf or ``False`` otherwise.
+
+ Returns:
+ A Python tree with the new values returned by ``fn``.
+ """
+
+def tree_map_with_path(
+ fn: Callable,
+ tree: Any,
+ *rest: Any,
+ 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.
+
+ This function is the same :func:`tree_map` but the ``fn`` takes the path as
+ the first argument followed by the remaining tree nodes.
+
+ Args:
+ fn (callable): The function that processes the leaves of the tree.
+ tree (Any): The main Python tree that will be iterated upon.
+ rest (tuple[Any]): Extra trees to be iterated together with ``tree``.
+ is_leaf (Optional[Callable]): An optional callable that returns ``True``
+ if the passed object is considered a leaf or ``False`` otherwise.
+ path (Optional[Any]): Prefix will be added to the result.
+
+ Returns:
+ A Python tree with the new values returned by ``fn``.
+
+ Example:
+ >>> from mlx.utils import tree_map_with_path
+ >>> tree = {"model": [{"w": 0, "b": 1}, {"w": 0, "b": 1}]}
+ >>> new_tree = tree_map_with_path(lambda path, _: print(path), tree)
+ model.0.w
+ model.0.b
+ model.1.w
+ model.1.b
+ """
+
+def tree_flatten(
+ tree: Any,
+ prefix: str = ...,
+ 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
+ complexity.
+
+ .. code-block:: python
+
+ from mlx.utils import tree_flatten
+
+ print(tree_flatten([[[0]]]))
+ # [("0.0.0", 0)]
+
+ print(tree_flatten([[[0]]], prefix=".hello"))
+ # [("hello.0.0.0", 0)]
+
+ tree_flatten({"a": {"b": 1}}, destination={})
+ {"a.b": 1}
+
+ .. note::
+ Dictionaries should have keys that are valid Python identifiers.
+
+ Args:
+ tree (Any): The Python tree to be flattened.
+ prefix (str): A prefix to use for the keys. The first character is
+ always discarded.
+ is_leaf (callable): An optional callable that returns True if the
+ passed object is considered a leaf or False otherwise.
+ destination (list or dict, optional): A list or dictionary to store the
+ flattened tree. If None an empty list will be used. Default: ``None``.
+
+ Returns:
+ Union[List[Tuple[str, Any]], Dict[str, Any]]: The flat representation of
+ the Python tree.
+ """
+
+def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
+ """Recreate a Python tree from its flat representation.
+
+ .. code-block:: python
+
+ from mlx.utils import tree_unflatten
+
+ d = tree_unflatten([("hello.world", 42)])
+ print(d)
+ # {"hello": {"world": 42}}
+
+ d = tree_unflatten({"hello.world": 42})
+ print(d)
+ # {"hello": {"world": 42}}
+
+ Args:
+ tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a Python tree.
+ For instance as returned by :meth:`tree_flatten`.
+
+ Returns:
+ A Python tree.
+ """
+
+def tree_reduce(fn, tree, initializer=..., is_leaf=...): # -> None:
+ """Applies a reduction to the leaves of a Python tree.
+
+ This function reduces Python trees into an accumulated result by applying
+ the provided function ``fn`` to the leaves of the tree.
+
+ Example:
+ >>> from mlx.utils import tree_reduce
+ >>> tree = {"a": [1, 2, 3], "b": [4, 5]}
+ >>> tree_reduce(lambda acc, x: acc + x, tree, 0)
+ 15
+
+ Args:
+ fn (callable): The reducer function that takes two arguments (accumulator,
+ current value) and returns the updated accumulator.
+ tree (Any): The Python tree to reduce. It can be any nested combination of
+ lists, tuples, or dictionaries.
+ initializer (Any, optional): The initial value to start the reduction. If
+ not provided, the first leaf value is used.
+ is_leaf (callable, optional): A function to determine if an object is a
+ leaf, returning ``True`` for leaf nodes and ``False`` otherwise.
+
+ Returns:
+ Any: The accumulated value.
+ """
+
+def tree_merge(
+ tree_a, tree_b, merge_fn=...
+): # -> dict[Any, Any] | list[Any] | tuple[Any, *tuple[Any, ...]] | tuple[Any, ...]:
+ """Merge two Python trees in one containing the values of both. It can be
+ thought of as a deep dict.update method.
+
+ Args:
+ tree_a (Any): The first Python tree.
+ tree_b (Any): The second Python tree.
+ merge_fn (callable, optional): A function to merge leaves.
+
+ Returns:
+ The Python tree containing the values of both ``tree_a`` and
+ ``tree_b``.
+ """
← 699fd959 fix exo scripts
·
back to Exo
·
show ips on dashboard e6068196 →