From 6bbb6344b64161a981869a2fd0f71eaeb31f5e2b Mon Sep 17 00:00:00 2001 From: rltakashige Date: Wed, 5 Nov 2025 21:26:04 -0800 Subject: [PATCH] mlx.distributed.Group type stubs --- .github/workflows/pipeline.yml | 27 + justfile | 2 +- pyproject.toml | 6 +- src/exo/engines/mlx/auto_parallel.py | 36 +- src/exo/engines/mlx/utils_mlx.py | 42 +- src/exo/worker/runner/generate.py | 14 +- src/exo/worker/runner/runner.py | 6 +- src/exo/worker/runner/runner_supervisor.py | 4 +- typings/mlx/core/__init__.pyi | 5420 ++++++++++++++++++++ typings/mlx/core/cuda/__init__.pyi | 2 + typings/mlx/core/distributed/__init__.pyi | 216 + typings/mlx/core/metal/__init__.pyi | 38 + typings/mlx/core/random/__init__.pyi | 301 ++ uv.lock | 20 +- 14 files changed, 6062 insertions(+), 72 deletions(-) create mode 100644 typings/mlx/core/__init__.pyi create mode 100644 typings/mlx/core/cuda/__init__.pyi create mode 100644 typings/mlx/core/distributed/__init__.pyi create mode 100644 typings/mlx/core/metal/__init__.pyi create mode 100644 typings/mlx/core/random/__init__.pyi diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml index 3fe6fa5b..25e240d4 100644 --- a/.github/workflows/pipeline.yml +++ b/.github/workflows/pipeline.yml @@ -62,6 +62,33 @@ jobs: fi shell: bash + - name: Configure basedpyright include for local MLX + run: | + RUNNER_LABELS='${{ toJSON(runner.labels) }}' + if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then + if [ -d "/Users/Shared/mlx" ]; then + echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx" + awk ' + BEGIN { in=0 } + /^\[tool\.basedpyright\]/ { in=1; print; next } + in && /^\[/ { in=0 } # next section + in && /^[ \t]*include[ \t]*=/ { + print "include = [\"/Users/Shared/mlx\"]" + next + } + { print } + ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml + + echo "New [tool.basedpyright] section:" + sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true + else + echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged." + fi + else + echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged." + fi + shell: bash + - uses: ./.github/actions/typecheck # ci: diff --git a/justfile b/justfile index 0db15c55..a61d0bb8 100644 --- a/justfile +++ b/justfile @@ -1,5 +1,5 @@ fmt: - uv run ruff format src + uv run ruff format src typings lint: uv run ruff check --fix src diff --git a/pyproject.toml b/pyproject.toml index d17ad793..6097e6ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ build-backend = "uv_build" ### [tool.basedpyright] +include = [".venv/lib/mlx", "src"] typeCheckingMode = "strict" failOnWarnings = true @@ -97,15 +98,12 @@ reportUnnecessaryTypeIgnoreComment = "error" pythonVersion = "3.13" pythonPlatform = "Darwin" -exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "mlx/*", "mlx-lm/*"] +exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust"] stubPath = "typings" [[tool.basedpyright.executionEnvironments]] root = "src" -[[tool.basedpyright.executionEnvironments]] -root = "." - ### # uv configuration ### diff --git a/src/exo/engines/mlx/auto_parallel.py b/src/exo/engines/mlx/auto_parallel.py index e5eee663..7db609d3 100644 --- a/src/exo/engines/mlx/auto_parallel.py +++ b/src/exo/engines/mlx/auto_parallel.py @@ -162,9 +162,9 @@ class PipelineParallelisationStrategy(ParallelisationShardStrategy): class TensorParallelisationStrategy(ParallelisationShardStrategy): - def __init__(self, group: mx.distributed.Group): # type: ignore - self.group = group # type: ignore - self.N = self.group.size # type: ignore + def __init__(self, group: mx.distributed.Group): + self.group = group + self.N = self.group.size def auto_parallel( self, model: nn.Module, model_shard_meta: ShardMetadata @@ -174,28 +174,28 @@ class TensorParallelisationStrategy(ParallelisationShardStrategy): all_to_sharded_linear = partial( shard_linear, sharding="all-to-sharded", - group=self.group, # pyright: ignore + group=self.group, ) sharded_to_all_linear = partial( shard_linear, sharding="sharded-to-all", - group=self.group, # type: ignore + group=self.group, ) all_to_sharded_linear_in_place = partial( shard_inplace, sharding="all-to-sharded", - group=self.group, # pyright: ignore + group=self.group, ) sharded_to_all_linear_in_place = partial( shard_inplace, sharding="sharded-to-all", - group=self.group, # type: ignore + group=self.group, ) if isinstance(model, LlamaModel): tensor_parallel_sharding_strategy = LlamaShardingStrategy( - self.group, # type: ignore + self.group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, @@ -203,7 +203,7 @@ class TensorParallelisationStrategy(ParallelisationShardStrategy): ) elif isinstance(model, DeepseekV3Model): tensor_parallel_sharding_strategy = DeepSeekShardingStrategy( - self.group, # type: ignore + self.group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, @@ -211,7 +211,7 @@ class TensorParallelisationStrategy(ParallelisationShardStrategy): ) elif isinstance(model, Qwen3MoeModel): tensor_parallel_sharding_strategy = QwenShardingStrategy( - self.group, # type: ignore + self.group, all_to_sharded_linear, sharded_to_all_linear, all_to_sharded_linear_in_place, @@ -305,14 +305,14 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy): class ShardedDeepseekV3MoE(CustomMlxLayer): def __init__(self, layer: _LayerCallable): super().__init__(layer) - self.sharding_group: mx.distributed.Group | None = None # type: ignore + self.sharding_group: mx.distributed.Group | None = None def __call__(self, x: mx.array) -> mx.array: - if self.sharding_group is not None: # type: ignore + if self.sharding_group is not None: x = sum_gradients(self.sharding_group)(x) # type: ignore y = self.original_layer.__call__(x) # type: ignore - if self.sharding_group is not None: # type: ignore - y = mx.distributed.all_sum(y, group=self.sharding_group) # type: ignore + if self.sharding_group is not None: + y = mx.distributed.all_sum(y, group=self.sharding_group) return y @@ -349,12 +349,12 @@ class QwenShardingStrategy(TensorParallelShardingStrategy): class ShardedQwenMoE(CustomMlxLayer): def __init__(self, layer: _LayerCallable): super().__init__(layer) - self.sharding_group: mx.distributed.Group | None = None # type: ignore + self.sharding_group: mx.distributed.Group | None = None def __call__(self, x: mx.array) -> mx.array: - if self.sharding_group is not None: # type: ignore + if self.sharding_group is not None: x = sum_gradients(self.sharding_group)(x) # type: ignore y = self.original_layer.__call__(x) # type: ignore - if self.sharding_group is not None: # type: ignore - y = mx.distributed.all_sum(y, group=self.sharding_group) # type: ignore + 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 5e730033..eb82246c 100644 --- a/src/exo/engines/mlx/utils_mlx.py +++ b/src/exo/engines/mlx/utils_mlx.py @@ -38,17 +38,17 @@ mlx_rank: None | int = None mlx_world_size: None | int = None -def mx_barrier(group: mx.distributed.Group | None = None): # type: ignore - mx.eval( # type: ignore +def mx_barrier(group: mx.distributed.Group | None = None): + mx.eval( mx.distributed.all_sum( mx.array(1.0), stream=mx.default_stream(mx.Device(mx.cpu)), - group=group, # type: ignore[type-arg] + group=group, ) ) -def broadcast_from_zero(value: int, group: mx.distributed.Group | None = None): # type: ignore +def broadcast_from_zero(value: int, group: mx.distributed.Group | None = None): if mlx_rank is None: return value @@ -57,8 +57,8 @@ def broadcast_from_zero(value: int, group: mx.distributed.Group | None = None): else: a = mx.array([0], dtype=mx.int32) - m = mx.distributed.all_sum(a, stream=mx.Device(mx.DeviceType.cpu), group=group) # type: ignore - mx.eval(m) # type: ignore + m = mx.distributed.all_sum(a, stream=mx.Device(mx.DeviceType.cpu), group=group) + mx.eval(m) return int(m.item()) @@ -68,12 +68,12 @@ class HostList(RootModel[list[str]]): return cls(root=[str(host) for host in hosts]) -def mlx_distributed_init( # type: ignore[return] +def mlx_distributed_init( rank: int, hosts: list[Host] | None = None, mlx_ibv_devices: list[list[str | None]] | None = None, mlx_ibv_coordinator: str | None = None, -) -> mx.distributed.Group: # type: ignore +) -> mx.distributed.Group: """ Initialize the MLX distributed (runs in thread pool). @@ -132,7 +132,9 @@ def initialize_mlx( hosts: list[Host] | None = None, mlx_ibv_devices: list[list[str | None]] | None = None, mlx_ibv_coordinator: str | None = None, -) -> tuple[Model, TokenizerWrapper, Callable[[mx.array], mx.array], Any]: +) -> tuple[ + Model, TokenizerWrapper, Callable[[mx.array], mx.array], mx.distributed.Group +]: """ Initialize the MLX model, tokenizer, and sampler. Runs in the MLX thread. @@ -141,7 +143,7 @@ def initialize_mlx( - mlx_ibv_devices: RDMA connectivity matrix """ mx.random.seed(42) - group = mlx_distributed_init( # type: ignore[misc] + group = mlx_distributed_init( model_shard_meta.device_rank, hosts=hosts, mlx_ibv_devices=mlx_ibv_devices, @@ -154,14 +156,14 @@ def initialize_mlx( sampler: Callable[[mx.array], mx.array] = make_sampler(temp=0.7) - model, tokenizer = shard_and_load(model_shard_meta, group=group) # type: ignore[reportUnknownArgumentType] + model, tokenizer = shard_and_load(model_shard_meta, group=group) return model, tokenizer, sampler, group # type: ignore[return-value] def shard_and_load( model_shard_meta: ShardMetadata, - group: mx.distributed.Group, # type: ignore + group: mx.distributed.Group, ) -> tuple[nn.Module, TokenizerWrapper]: model_path = build_model_path(model_shard_meta.model_meta.model_id) @@ -177,7 +179,7 @@ def shard_and_load( assert isinstance(tokenizer, _TokenizerWrapper) if group: - runner_print(f"Group size: {group.size()}, group rank: {group.rank()}") # type: ignore + runner_print(f"Group size: {group.size()}, group rank: {group.rank()}") else: runner_print("!!! No group") @@ -189,19 +191,19 @@ def shard_and_load( case "pipeline_rdma": strategy = PipelineParallelisationStrategy() case "tensor": - strategy = TensorParallelisationStrategy(group) # type: ignore[reportUnknownArgumentType] + strategy = TensorParallelisationStrategy(group) case "tensor_rdma": - strategy = TensorParallelisationStrategy(group) # type: ignore[reportUnknownArgumentType] + strategy = TensorParallelisationStrategy(group) model = strategy.auto_parallel(model, model_shard_meta) runner_print(f"Model after auto_parallel: {str(model)}") mx.eval(model.parameters()) # type: ignore - mx.eval(model) # type: ignore + mx.eval(model) # Synchronize processes before generation to avoid timeout - mx_barrier(group) # type: ignore[reportUnknownArgumentType] + mx_barrier(group) return model, tokenizer # type: ignore @@ -288,15 +290,15 @@ def mlx_force_oom(size: int = 40000) -> None: """ Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations. """ - mx.set_default_device(mx.gpu) # type: ignore + mx.set_default_device(mx.gpu) a = mx.random.uniform(shape=(size, size), dtype=mx.float32) b = mx.random.uniform(shape=(size, size), dtype=mx.float32) - mx.eval(a, b) # type: ignore + mx.eval(a, b) c = mx.matmul(a, b) d = mx.matmul(a, c) e = mx.matmul(b, c) f = mx.sigmoid(d + e) - mx.eval(f) # type: ignore + mx.eval(f) def set_wired_limit_for_model(model_size: Memory): diff --git a/src/exo/worker/runner/generate.py b/src/exo/worker/runner/generate.py index eb207842..3db14141 100644 --- a/src/exo/worker/runner/generate.py +++ b/src/exo/worker/runner/generate.py @@ -14,9 +14,9 @@ from mlx_lm.models.cache import KVCache from exo.engines.mlx import Model, TokenizerWrapper from exo.engines.mlx.utils_mlx import ( apply_chat_template, - broadcast_from_zero, # type: ignore + broadcast_from_zero, make_kv_cache, - mx_barrier, # type: ignore + mx_barrier, ) from exo.shared.types.api import ChatCompletionMessage from exo.shared.types.tasks import ChatCompletionTaskParams @@ -62,7 +62,7 @@ def generate_step( quantized_kv_start: int = 0, prompt_progress_callback: Callable[[int, int], None] | None = None, input_embeddings: mx.array | None = None, - group: mx.distributed.Group | None = None, # type: ignore[type-arg] + group: mx.distributed.Group | None = None, ) -> Generator[Tuple[int, mx.array], None, None]: """ A generator producing token ids based on the given prompt from the model. @@ -213,7 +213,7 @@ def generate_step( y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings) - mx.async_eval(y, logprobs) # type: ignore[type-arg] + mx.async_eval(y, logprobs) next_y: array | None = None next_logprobs: array | None = None n = 0 @@ -221,7 +221,7 @@ def generate_step( if n != max_tokens: assert y is not None next_y, next_logprobs = _step(y) - mx.async_eval(next_y, next_logprobs) # type: ignore[type-arg] + 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) @@ -250,7 +250,7 @@ def stream_generate( quantized_kv_start: int = 0, prompt_progress_callback: Callable[[int, int], None] | None = None, input_embeddings: mx.array | None = None, - group: mx.distributed.Group | None = None, # type: ignore[type-arg] + group: mx.distributed.Group | None = None, ) -> Generator[GenerationResponse, None, None]: # Try to infer if special tokens are needed add_special_tokens = tokenizer.bos_token is None or not prompt.startswith( @@ -310,7 +310,7 @@ async def warmup_inference( model: Model, tokenizer: TokenizerWrapper, sampler: Callable[[mx.array], mx.array], - group: mx.distributed.Group | None = None, # type: ignore + group: mx.distributed.Group | None = None, ) -> int: loop = asyncio.get_running_loop() diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index f7fe305a..78b782da 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -25,7 +25,7 @@ from exo.shared.types.worker.communication import ( ) from exo.shared.types.worker.shards import ShardMetadata from exo.utils import ensure_type -from exo.worker.runner.generate import mlx_generate, warmup_inference # type: ignore +from exo.worker.runner.generate import mlx_generate, warmup_inference async def main(raw_conn: Connection): @@ -51,7 +51,7 @@ async def main(raw_conn: Connection): mlx_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) loop = asyncio.get_running_loop() - model, tokenizer, sampler, group = await loop.run_in_executor( # type: ignore[type-arg] + model, tokenizer, sampler, group = await loop.run_in_executor( mlx_executor, partial( initialize_mlx, @@ -70,7 +70,7 @@ async def main(raw_conn: Connection): model=model, tokenizer=tokenizer, sampler=sampler, - group=group, # type: ignore[type-arg] + group=group, ) runner_print(f"Warmed up by generating {toks} tokens") await conn.send(InitializedResponse(time_taken=time.time() - setup_start_time)) diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 63efbe88..1923ac96 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -221,9 +221,7 @@ class RunnerSupervisor: timeout = PREFILL_TIMEOUT_SECONDS - logger.info( - f"Starting chat completion with timeout {timeout}" - ) + logger.info(f"Starting chat completion with timeout {timeout}") while True: try: diff --git a/typings/mlx/core/__init__.pyi b/typings/mlx/core/__init__.pyi new file mode 100644 index 00000000..e1ffbe29 --- /dev/null +++ b/typings/mlx/core/__init__.pyi @@ -0,0 +1,5420 @@ +import enum +import pathlib +import sys +import types +from typing import ( + Annotated, + Any, + Callable, + Literal, + Mapping, + Sequence, + TypeAlias, + overload, +) + +import numpy +from numpy.typing import ArrayLike as _ArrayLike + +from . import cuda as cuda +from . import distributed as distributed +from . import metal as metal +from . import random as random + +class ArrayAt: + """A helper object to apply updates at specific indices.""" + def __getitem__(self, indices: object | None) -> ArrayAt: ... + def add( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def subtract( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def multiply( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def divide( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def maximum( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def minimum( + self, + value: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + +class ArrayIterator: + """A helper object to iterate over the 1st dimension of an array.""" + def __next__(self) -> array: ... + def __iter__(self) -> ArrayIterator: ... + +class ArrayLike: + """ + Any Python object which has an ``__mlx__array__`` method that + returns an :obj:`array`. + """ + def __init__(self, arg: object, /) -> None: ... + +class Device: + """A device to run operations on.""" + def __init__(self, type: DeviceType, index: int = ...) -> None: ... + @property + def type(self) -> DeviceType: ... + def __repr__(self) -> str: ... + def __eq__(self, arg: object, /) -> bool: ... + +class DeviceType(enum.Enum): + cpu = ... # type: ignore + gpu = ... #  type: ignore + def __eq__(self, arg: object, /) -> bool: ... + +class Dtype: + """ + An object to hold the type of a :class:`array`. + + See the :ref:`list of types ` for more details + on available data types. + """ + @property + def size(self) -> int: + """Size of the type in bytes.""" + + def __repr__(self) -> str: ... + def __eq__(self, arg: object, /) -> bool: ... + def __hash__(self) -> int: ... + +class DtypeCategory(enum.Enum): + """ + Type to hold categories of :class:`dtypes `. + + * :attr:`~mlx.core.generic` + + * :ref:`bool_ ` + * :attr:`~mlx.core.number` + + * :attr:`~mlx.core.integer` + + * :attr:`~mlx.core.unsignedinteger` + + * :ref:`uint8 ` + * :ref:`uint16 ` + * :ref:`uint32 ` + * :ref:`uint64 ` + + * :attr:`~mlx.core.signedinteger` + + * :ref:`int8 ` + * :ref:`int32 ` + * :ref:`int64 ` + + * :attr:`~mlx.core.inexact` + + * :attr:`~mlx.core.floating` + + * :ref:`float16 ` + * :ref:`bfloat16 ` + * :ref:`float32 ` + * :ref:`float64 ` + + * :attr:`~mlx.core.complexfloating` + + * :ref:`complex64 ` + + See also :func:`~mlx.core.issubdtype`. + """ + + complexfloating = ... + floating = ... + inexact = ... + signedinteger = ... + unsignedinteger = ... + integer = ... + number = ... + generic = ... + +class FunctionExporter: + """ + A context managing class for exporting multiple traces of the same + function to a file. + + Make an instance of this class by calling fun:`mx.exporter`. + """ + def close(self) -> None: ... + def __enter__(self) -> FunctionExporter: ... + def __exit__( + self, + exc_type: object | None = ..., + exc_value: object | None = ..., + traceback: object | None = ..., + ) -> None: ... + def __call__(self, *args, **kwargs) -> None: ... + +class Stream: + """A stream for running operations on a given device.""" + @property + def device(self) -> Device: ... + def __repr__(self) -> str: ... + def __eq__(self, arg: object, /) -> bool: ... + +class StreamContext: + """ + A context manager for setting the current device and stream. + + See :func:`stream` for usage. + + Args: + s: The stream or device to set as the default. + """ + def __init__(self, s: Stream | Device) -> None: ... + def __enter__(self) -> None: ... + def __exit__( + self, + exc_type: type | None = ..., + exc_value: object | None = ..., + traceback: object | None = ..., + ) -> None: ... + +def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise absolute value. + + Args: + a (array): Input array. + + Returns: + array: The absolute value of ``a``. + """ + +def add( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise addition. + + Add two arrays with numpy-style broadcasting semantics. Either or both input arrays + can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The sum of ``a`` and ``b``. + """ + +def addmm( + c: array, + a: array, + b: array, + /, + alpha: float = ..., + beta: float = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Matrix multiplication with addition and optional scaling. + + Perform the (possibly batched) matrix multiplication of two arrays and add to the result + with optional scaling factors. + + Args: + c (array): Input array or scalar. + a (array): Input array or scalar. + b (array): Input array or scalar. + alpha (float, optional): Scaling factor for the + matrix product of ``a`` and ``b`` (default: ``1``) + beta (float, optional): Scaling factor for ``c`` (default: ``1``) + + Returns: + array: ``alpha * (a @ b) + beta * c`` + """ + +def all( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + An `and` reduction over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def allclose( + a: array, + b: array, + /, + rtol: float = ..., + atol: float = ..., + *, + equal_nan: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Approximate comparison of two arrays. + + Infinite values are considered equal if they have the same sign, NaN values are not equal unless ``equal_nan`` is ``True``. + + The arrays are considered equal if: + + .. code-block:: + + all(abs(a - b) <= (atol + rtol * abs(b))) + + Note unlike :func:`array_equal`, this function supports numpy-style + broadcasting. + + Args: + a (array): Input array. + b (array): Input array. + rtol (float): Relative tolerance. + atol (float): Absolute tolerance. + equal_nan (bool): If ``True``, NaNs are considered equal. + Defaults to ``False``. + + Returns: + array: The boolean output scalar indicating if the arrays are close. + """ + +def any( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + An `or` reduction over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +@overload +def arange( + start: int | float, + stop: int | float, + step: int | float | None, + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Generates ranges of numbers. + + Generate numbers in the half-open interval ``[start, stop)`` in + increments of ``step``. + + Args: + start (float or int, optional): Starting value which defaults to ``0``. + stop (float or int): Stopping value. + step (float or int, optional): Increment which defaults to ``1``. + dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``. + + Returns: + array: The range of values. + + Note: + Following the Numpy convention the actual increment used to + generate numbers is ``dtype(start + step) - dtype(start)``. + This can lead to unexpected results for example if `start + step` + is a fractional value and the `dtype` is integral. + """ + +@overload +def arange( + stop: int | float, + step: int | float | None = ..., + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: ... +def arccos(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse cosine. + + Args: + a (array): Input array. + + Returns: + array: The inverse cosine of ``a``. + """ + +def arccosh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse hyperbolic cosine. + + Args: + a (array): Input array. + + Returns: + array: The inverse hyperbolic cosine of ``a``. + """ + +def arcsin(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse sine. + + Args: + a (array): Input array. + + Returns: + array: The inverse sine of ``a``. + """ + +def arcsinh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse hyperbolic sine. + + Args: + a (array): Input array. + + Returns: + array: The inverse hyperbolic sine of ``a``. + """ + +def arctan(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse tangent. + + Args: + a (array): Input array. + + Returns: + array: The inverse tangent of ``a``. + """ + +def arctan2(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse tangent of the ratio of two arrays. + + Args: + a (array): Input array. + b (array): Input array. + + Returns: + array: The inverse tangent of the ratio of ``a`` and ``b``. + """ + +def arctanh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse hyperbolic tangent. + + Args: + a (array): Input array. + + Returns: + array: The inverse hyperbolic tangent of ``a``. + """ + +def argmax( + a: array, + /, + axis: int | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Indices of the maximum values along the axis. + + Args: + a (array): Input array. + axis (int, optional): Optional axis to reduce over. If unspecified + this defaults to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The ``uint32`` array with the indices of the maximum values. + """ + +def argmin( + a: array, + /, + axis: int | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Indices of the minimum values along the axis. + + Args: + a (array): Input array. + axis (int, optional): Optional axis to reduce over. If unspecified + this defaults to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The ``uint32`` array with the indices of the minimum values. + """ + +def argpartition( + a: array, + /, + kth: int, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Returns the indices that partition the array. + + The ordering of the elements within a partition in given by the indices + is undefined. + + Args: + a (array): Input array. + kth (int): Element index at the ``kth`` position in the output will + give the sorted position. All indices before the ``kth`` position + will be of elements less or equal to the element at the ``kth`` + index and all indices after will be of elements greater or equal + to the element at the ``kth`` index. + axis (int or None, optional): Optional axis to partition over. + If ``None``, this partitions over the flattened array. + If unspecified, it defaults to ``-1``. + + Returns: + array: The ``uint32`` array containing indices that partition the input. + """ + +def argsort( + a: array, + /, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Returns the indices that sort the array. + + Args: + a (array): Input array. + axis (int or None, optional): Optional axis to sort over. + If ``None``, this sorts over the flattened array. + If unspecified, it defaults to -1 (sorting over the last axis). + + Returns: + array: The ``uint32`` array containing indices that sort the input. + """ + +class array: + """An N-dimensional array object.""" + def __init__( + self: array, + val: scalar | list | tuple | numpy.ndarray | array, + dtype: Dtype | None = ..., + ) -> None: ... + def __buffer__(self, flags, /): + """ + Return a buffer object that exposes the underlying memory of the object. + """ + + def __release_buffer__(self, buffer, /): + """ + Release the buffer object that exposes the underlying memory of the object. + """ + + @property + def size(self) -> int: + """Number of elements in the array.""" + + @property + def ndim(self) -> int: + """The array's dimension.""" + + @property + def itemsize(self) -> int: + """The size of the array's datatype in bytes.""" + + @property + def nbytes(self) -> int: + """The number of bytes in the array.""" + + @property + def shape(self) -> tuple[int, ...]: + """ + The shape of the array as a Python tuple. + + Returns: + tuple(int): A tuple containing the sizes of each dimension. + """ + + @property + def dtype(self) -> Dtype: + """The array's :class:`Dtype`.""" + + @property + def real(self) -> array: + """The real part of a complex array.""" + + @property + def imag(self) -> array: + """The imaginary part of a complex array.""" + + def item(self) -> scalar: + """ + Access the value of a scalar array. + + Returns: + Standard Python scalar. + """ + + def tolist(self) -> list_or_scalar: + """ + Convert the array to a Python :class:`list`. + + Returns: + list: The Python list. + + If the array is a scalar then a standard Python scalar is returned. + + If the array has more than one dimension then the result is a nested + list of lists. + + The value type of the list corresponding to the last dimension is either + ``bool``, ``int`` or ``float`` depending on the ``dtype`` of the array. + """ + + def astype(self, dtype: Dtype, stream: Stream | Device | None = ...) -> array: + """ + Cast the array to a specified type. + + Args: + dtype (Dtype): Type to which the array is cast. + stream (Stream): Stream (or device) for the operation. + + Returns: + array: The array with type ``dtype``. + """ + + def __array_namespace__(self, api_version: str | None = ...) -> types.ModuleType: + """ + Returns an object that has all the array API functions on it. + + See the `Python array API `_ + for more information. + + Args: + api_version (str, optional): String representing the version + of the array API spec to return. Default: ``None``. + + Returns: + out (Any): An object representing the array API namespace. + """ + + def __getitem__(self, arg: object | None) -> array: ... + def __setitem__( + self, + arg0: object | None, + arg1: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> None: ... + @property + def at(self) -> ArrayAt: + """ + Used to apply updates at the given indices. + + .. note:: + + Regular in-place updates map to assignment. For instance ``x[idx] += y`` + maps to ``x[idx] = x[idx] + y``. As a result, assigning to the + same index ignores all but one update. Using ``x.at[idx].add(y)`` + will correctly apply all updates to all indices. + + .. list-table:: + :header-rows: 1 + + * - array.at syntax + - In-place syntax + * - ``x = x.at[idx].add(y)`` + - ``x[idx] += y`` + * - ``x = x.at[idx].subtract(y)`` + - ``x[idx] -= y`` + * - ``x = x.at[idx].multiply(y)`` + - ``x[idx] *= y`` + * - ``x = x.at[idx].divide(y)`` + - ``x[idx] /= y`` + * - ``x = x.at[idx].maximum(y)`` + - ``x[idx] = mx.maximum(x[idx], y)`` + * - ``x = x.at[idx].minimum(y)`` + - ``x[idx] = mx.minimum(x[idx], y)`` + + Example: + >>> a = mx.array([0, 0]) + >>> idx = mx.array([0, 1, 0, 1]) + >>> a[idx] += 1 + >>> a + array([1, 1], dtype=int32) + >>> + >>> a = mx.array([0, 0]) + >>> a.at[idx].add(1) + array([2, 2], dtype=int32) + """ + + def __len__(self) -> int: ... + def __iter__(self) -> ArrayIterator: ... + def __getstate__(self) -> tuple: ... + def __setstate__(self, arg: tuple, /) -> None: ... + def __dlpack__(self) -> _ArrayLike: ... + def __dlpack_device__(self) -> tuple: ... + def __copy__(self) -> array: ... + def __deepcopy__(self, memo: dict) -> array: ... + def __add__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __iadd__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __radd__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __sub__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __isub__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rsub__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __mul__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __imul__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rmul__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __truediv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __itruediv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rtruediv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __div__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rdiv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __floordiv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ifloordiv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rfloordiv__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __mod__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __imod__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rmod__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __eq__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array | bool: ... + def __lt__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __le__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __gt__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ge__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ne__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array | bool: ... + def __neg__(self) -> array: ... + def __bool__(self) -> bool: ... + def __repr__(self) -> str: ... + def __matmul__(self, other: array) -> array: ... + def __imatmul__(self, other: array) -> array: ... + def __pow__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rpow__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ipow__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __invert__(self) -> array: ... + def __and__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __iand__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __or__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ior__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __lshift__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ilshift__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __rshift__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __irshift__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __xor__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __ixor__( + self, + other: bool + | int + | float + | array + | Annotated[_ArrayLike, dict(order="C", device="cpu", writable=False)] + | complex + | ArrayLike, + ) -> array: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def flatten( + self, + start_axis: int = ..., + end_axis: int = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`flatten`.""" + + def reshape(self, *shape, stream: Stream | Device | None = ...) -> array: + """ + Equivalent to :func:`reshape` but the shape can be passed either as a + :obj:`tuple` or as separate arguments. + + See :func:`reshape` for full documentation. + """ + + def squeeze( + self, + axis: int | Sequence[int] | None = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`squeeze`.""" + + def abs(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`abs`.""" + + def __abs__(self) -> array: + """See :func:`abs`.""" + + def square(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`square`.""" + + def sqrt(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`sqrt`.""" + + def rsqrt(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`rsqrt`.""" + + def reciprocal(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`reciprocal`.""" + + def exp(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`exp`.""" + + def log(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`log`.""" + + def log2(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`log2`.""" + + def log10(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`log10`.""" + + def sin(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`sin`.""" + + def cos(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`cos`.""" + + def log1p(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`log1p`.""" + + def all( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`all`.""" + + def any( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`any`.""" + + def moveaxis( + self, source: int, destination: int, *, stream: Stream | Device | None = ... + ) -> array: + """See :func:`moveaxis`.""" + + def swapaxes( + self, axis1: int, axis2: int, *, stream: Stream | Device | None = ... + ) -> array: + """See :func:`swapaxes`.""" + + def transpose(self, *axes, stream: Stream | Device | None = ...) -> array: + """ + Equivalent to :func:`transpose` but the axes can be passed either as + a tuple or as separate arguments. + + See :func:`transpose` for full documentation. + """ + + @property + def T(self) -> array: + """Equivalent to calling ``self.transpose()`` with no arguments.""" + + def sum( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`sum`.""" + + def prod( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`prod`.""" + + def min( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`min`.""" + + def max( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`max`.""" + + def logcumsumexp( + self, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`logcumsumexp`.""" + + def logsumexp( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`logsumexp`.""" + + def mean( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`mean`.""" + + def std( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + ddof: int = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`std`.""" + + def var( + self, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + ddof: int = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`var`.""" + + def split( + self, + indices_or_sections: int | tuple[int, ...], + axis: int = ..., + *, + stream: Stream | Device | None = ..., + ) -> list[array]: + """See :func:`split`.""" + + def argmin( + self, + axis: int | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`argmin`.""" + + def argmax( + self, + axis: int | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`argmax`.""" + + def cumsum( + self, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`cumsum`.""" + + def cumprod( + self, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`cumprod`.""" + + def cummax( + self, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`cummax`.""" + + def cummin( + self, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`cummin`.""" + + def round( + self, decimals: int = ..., *, stream: Stream | Device | None = ... + ) -> array: + """See :func:`round`.""" + + def diagonal( + self, + offset: int = ..., + axis1: int = ..., + axis2: int = ..., + stream: Stream | Device | None = ..., + ) -> array: + """See :func:`diagonal`.""" + + def diag(self, k: int = ..., *, stream: Stream | Device | None = ...) -> array: + """Extract a diagonal or construct a diagonal matrix.""" + + def conj(self, *, stream: Stream | Device | None = ...) -> array: + """See :func:`conj`.""" + + def view(self, dtype: Dtype, *, stream: Stream | Device | None = ...) -> array: + """See :func:`view`.""" + +def array_equal( + a: scalar | array, + b: scalar | array, + equal_nan: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Array equality check. + + Compare two arrays for equality. Returns ``True`` if and only if the arrays + have the same shape and their values are equal. The arrays need not have + the same type to be considered equal. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + equal_nan (bool): If ``True``, NaNs are considered equal. + Defaults to ``False``. + + Returns: + array: A scalar boolean array. + """ + +def as_strided( + a: array, + /, + shape: Sequence[int] | None = ..., + strides: Sequence[int] | None = ..., + offset: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Create a view into the array with the given shape and strides. + + The resulting array will always be as if the provided array was row + contiguous regardless of the provided arrays storage order and current + strides. + + .. note:: + Note that this function should be used with caution as it changes + the shape and strides of the array directly. This can lead to the + resulting array pointing to invalid memory locations which can + result into crashes. + + Args: + a (array): Input array + shape (list(int), optional): The shape of the resulting array. If + None it defaults to ``a.shape()``. + strides (list(int), optional): The strides of the resulting array. If + None it defaults to the reverse exclusive cumulative product of + ``a.shape()``. + offset (int): Skip that many elements from the beginning of the input + array. + + Returns: + array: The output array which is the strided view of the input. + """ + +def async_eval(*args: MX_ARRAY_TREE) -> None: + """ + Asynchronously evaluate an :class:`array` or tree of :class:`array`. + + .. note:: + + This is an experimental API and may change in future versions. + + Args: + *args (arrays or trees of arrays): Each argument can be a single array + or a tree of arrays. If a tree is given the nodes can be a Python + :class:`list`, :class:`tuple` or :class:`dict`. Leaves which are not + arrays are ignored. + + Example: + >>> x = mx.array(1.0) + >>> y = mx.exp(x) + >>> mx.async_eval(y) + >>> print(y) + >>> + >>> y = mx.exp(x) + >>> mx.async_eval(y) + >>> z = y + 3 + >>> mx.async_eval(z) + >>> print(z) + """ + +def atleast_1d( + *arys: array, stream: Stream | Device | None = ... +) -> array | list[array]: + """ + Convert all arrays to have at least one dimension. + + Args: + *arys: Input arrays. + stream (Stream | Device | None, optional): The stream to execute the operation on. + + Returns: + array or list(array): An array or list of arrays with at least one dimension. + """ + +def atleast_2d( + *arys: array, stream: Stream | Device | None = ... +) -> array | list[array]: + """ + Convert all arrays to have at least two dimensions. + + Args: + *arys: Input arrays. + stream (Stream | Device | None, optional): The stream to execute the operation on. + + Returns: + array or list(array): An array or list of arrays with at least two dimensions. + """ + +def atleast_3d( + *arys: array, stream: Stream | Device | None = ... +) -> array | list[array]: + """ + Convert all arrays to have at least three dimensions. + + Args: + *arys: Input arrays. + stream (Stream | Device | None, optional): The stream to execute the operation on. + + Returns: + array or list(array): An array or list of arrays with at least three dimensions. + """ + +bfloat16: Dtype = ... + +def bitwise_and( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise bitwise and. + + Take the bitwise and of two arrays with numpy-style broadcasting + semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The bitwise and ``a & b``. + """ + +def bitwise_invert(a: scalar | array, stream: Stream | Device | None = ...) -> array: + """ + Element-wise bitwise inverse. + + Take the bitwise complement of the input. + + Args: + a (array): Input array or scalar. + + Returns: + array: The bitwise inverse ``~a``. + """ + +def bitwise_or( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise bitwise or. + + Take the bitwise or of two arrays with numpy-style broadcasting + semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The bitwise or``a | b``. + """ + +def bitwise_xor( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise bitwise xor. + + Take the bitwise exclusive or of two arrays with numpy-style + broadcasting semantics. Either or both input arrays can also be + scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The bitwise xor ``a ^ b``. + """ + +def block_masked_mm( + a: array, + b: array, + /, + block_size: int = ..., + mask_out: array | None = ..., + mask_lhs: array | None = ..., + mask_rhs: array | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + r""" + Matrix multiplication with block masking. + + Perform the (possibly batched) matrix multiplication of two arrays and with blocks + of size ``block_size x block_size`` optionally masked out. + + Assuming ``a`` with shape (..., `M`, `K`) and b with shape (..., `K`, `N`) + + * ``lhs_mask`` must have shape (..., :math:`\lceil` `M` / ``block_size`` :math:`\rceil`, :math:`\lceil` `K` / ``block_size`` :math:`\rceil`) + + * ``rhs_mask`` must have shape (..., :math:`\lceil` `K` / ``block_size`` :math:`\rceil`, :math:`\lceil` `N` / ``block_size`` :math:`\rceil`) + + * ``out_mask`` must have shape (..., :math:`\lceil` `M` / ``block_size`` :math:`\rceil`, :math:`\lceil` `N` / ``block_size`` :math:`\rceil`) + + Note: Only ``block_size=64`` and ``block_size=32`` are currently supported + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + block_size (int): Size of blocks to be masked. Must be ``32`` or ``64``. Default: ``64``. + mask_out (array, optional): Mask for output. Default: ``None``. + mask_lhs (array, optional): Mask for ``a``. Default: ``None``. + mask_rhs (array, optional): Mask for ``b``. Default: ``None``. + + Returns: + array: The output array. + """ + +def broadcast_arrays( + *arrays: array, stream: Stream | Device | None = ... +) -> tuple[array, ...]: + """ + Broadcast arrays against one another. + + The broadcasting semantics are the same as Numpy. + + Args: + *arrays (array): The input arrays. + + Returns: + tuple(array): The output arrays with the broadcasted shape. + """ + +def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int]: + """ + Broadcast shapes. + + Returns the shape that results from broadcasting the supplied array shapes + against each other. + + Args: + *shapes (Sequence[int]): The shapes to broadcast. + + Returns: + tuple: The broadcasted shape. + + Raises: + ValueError: If the shapes cannot be broadcast. + + Example: + >>> mx.broadcast_shapes((1,), (3, 1)) + (3, 1) + >>> mx.broadcast_shapes((6, 7), (5, 6, 1), (7,)) + (5, 6, 7) + >>> mx.broadcast_shapes((5, 1, 4), (1, 3, 1)) + (5, 3, 4) + """ + +def broadcast_to( + a: scalar | array, + /, + shape: Sequence[int], + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Broadcast an array to the given shape. + + The broadcasting semantics are the same as Numpy. + + Args: + a (array): Input array. + shape (list(int)): The shape to broadcast to. + + Returns: + array: The output array with the new shape. + """ + +def ceil(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise ceil. + + Args: + a (array): Input array. + + Returns: + array: The ceil of ``a``. + """ + +def checkpoint(fun: Callable) -> Callable: ... +def clear_cache() -> None: + """ + Clear the memory cache. + + After calling this, :func:`get_cache_memory` should return ``0``. + """ + +def clip( + a: array, + /, + a_min: scalar | array | None, + a_max: scalar | array | None, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Clip the values of the array between the given minimum and maximum. + + If either ``a_min`` or ``a_max`` are ``None``, then corresponding edge + is ignored. At least one of ``a_min`` and ``a_max`` cannot be ``None``. + The input ``a`` and the limits must broadcast with one another. + + Args: + a (array): Input array. + a_min (scalar or array or None): Minimum value to clip to. + a_max (scalar or array or None): Maximum value to clip to. + + Returns: + array: The clipped array. + """ + +def compile( + fun: Callable, + inputs: object | None = ..., + outputs: object | None = ..., + shapeless: bool = ..., +) -> Callable: + """ + Returns a compiled function which produces the same output as ``fun``. + + Args: + fun (Callable): A function which takes a variable number of + :class:`array` or trees of :class:`array` and returns + a variable number of :class:`array` or trees of :class:`array`. + inputs (list or dict, optional): These inputs will be captured during + the function compilation along with the inputs to ``fun``. The ``inputs`` + can be a :obj:`list` or a :obj:`dict` containing arbitrarily nested + lists, dictionaries, or arrays. Leaf nodes that are not + :obj:`array` are ignored. Default: ``None`` + outputs (list or dict, optional): These outputs will be captured and + updated in a compiled function. The ``outputs`` can be a + :obj:`list` or a :obj:`dict` containing arbitrarily nested lists, + dictionaries, or arrays. Leaf nodes that are not :obj:`array` are ignored. + Default: ``None`` + shapeless (bool, optional): A function compiled with the ``shapeless`` + option enabled will not be recompiled when the input shape changes. Not all + functions can be compiled with ``shapeless`` enabled. Attempting to compile + such functions with shapeless enabled will throw. Note, changing the number + of dimensions or type of any input will result in a recompilation even with + ``shapeless`` set to ``True``. Default: ``False`` + + Returns: + Callable: A compiled function which has the same input arguments + as ``fun`` and returns the the same output(s). + """ + +complex64: Dtype = ... +complexfloating: DtypeCategory = ... + +def concat( + arrays: list[array], + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """See :func:`concatenate`.""" + +def concatenate( + arrays: list[array], + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Concatenate the arrays along the given axis. + + Args: + arrays (list(array)): Input :obj:`list` or :obj:`tuple` of arrays. + axis (int, optional): Optional axis to concatenate along. If + unspecified defaults to ``0``. + + Returns: + array: The concatenated array. + """ + +def conj(a: array, *, stream: Stream | Device | None = ...) -> array: + """ + Return the elementwise complex conjugate of the input. + Alias for `mx.conjugate`. + + Args: + a (array): Input array + + Returns: + array: The output array. + """ + +def conjugate(a: array, *, stream: Stream | Device | None = ...) -> array: + """ + Return the elementwise complex conjugate of the input. + Alias for `mx.conj`. + + Args: + a (array): Input array + + Returns: + array: The output array. + """ + +def contiguous( + a: array, + /, + allow_col_major: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Force an array to be row contiguous. Copy if necessary. + + Args: + a (array): The input to make contiguous + allow_col_major (bool): Consider column major as contiguous and don't copy + + Returns: + array: The row or col contiguous output. + """ + +def conv1d( + input: array, + weight: array, + /, + stride: int = ..., + padding: int = ..., + dilation: int = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 1D convolution over an input with several channels + + Args: + input (array): Input array of shape ``(N, L, C_in)``. + weight (array): Weight array of shape ``(C_out, K, C_in)``. + stride (int, optional): Kernel stride. Default: ``1``. + padding (int, optional): Input padding. Default: ``0``. + dilation (int, optional): Kernel dilation. Default: ``1``. + groups (int, optional): Input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def conv2d( + input: array, + weight: array, + /, + stride: int | tuple[int, int] = ..., + padding: int | tuple[int, int] = ..., + dilation: int | tuple[int, int] = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 2D convolution over an input with several channels + + Args: + input (array): Input array of shape ``(N, H, W, C_in)``. + weight (array): Weight array of shape ``(C_out, KH, KW, C_in)``. + stride (int or tuple(int), optional): :obj:`tuple` of size 2 with + kernel strides. All spatial dimensions get the same stride if + only one number is specified. Default: ``1``. + padding (int or tuple(int), optional): :obj:`tuple` of size 2 with + symmetric input padding. All spatial dimensions get the same + padding if only one number is specified. Default: ``0``. + dilation (int or tuple(int), optional): :obj:`tuple` of size 2 with + kernel dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + groups (int, optional): input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def conv3d( + input: array, + weight: array, + /, + stride: int | tuple[int, int, int] = ..., + padding: int | tuple[int, int, int] = ..., + dilation: int | tuple[int, int, int] = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 3D convolution over an input with several channels + + Note: Only the default ``groups=1`` is currently supported. + + Args: + input (array): Input array of shape ``(N, D, H, W, C_in)``. + weight (array): Weight array of shape ``(C_out, KD, KH, KW, C_in)``. + stride (int or tuple(int), optional): :obj:`tuple` of size 3 with + kernel strides. All spatial dimensions get the same stride if + only one number is specified. Default: ``1``. + padding (int or tuple(int), optional): :obj:`tuple` of size 3 with + symmetric input padding. All spatial dimensions get the same + padding if only one number is specified. Default: ``0``. + dilation (int or tuple(int), optional): :obj:`tuple` of size 3 with + kernel dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + groups (int, optional): input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def conv_general( + input: array, + weight: array, + /, + stride: int | Sequence[int] = ..., + padding: int | Sequence[int] | tuple[Sequence[int] | Sequence[int]] = ..., + kernel_dilation: int | Sequence[int] = ..., + input_dilation: int | Sequence[int] = ..., + groups: int = ..., + flip: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + General convolution over an input with several channels + + Args: + input (array): Input array of shape ``(N, ..., C_in)``. + weight (array): Weight array of shape ``(C_out, ..., C_in)``. + stride (int or list(int), optional): :obj:`list` with kernel strides. + All spatial dimensions get the same stride if + only one number is specified. Default: ``1``. + padding (int, list(int), or tuple(list(int), list(int)), optional): + :obj:`list` with input padding. All spatial dimensions get the same + padding if only one number is specified. Default: ``0``. + kernel_dilation (int or list(int), optional): :obj:`list` with + kernel dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + input_dilation (int or list(int), optional): :obj:`list` with + input dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + groups (int, optional): Input feature groups. Default: ``1``. + flip (bool, optional): Flip the order in which the spatial dimensions of + the weights are processed. Performs the cross-correlation operator when + ``flip`` is ``False`` and the convolution operator otherwise. + Default: ``False``. + + Returns: + array: The convolved array. + """ + +def conv_transpose1d( + input: array, + weight: array, + /, + stride: int = ..., + padding: int = ..., + dilation: int = ..., + output_padding: int = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 1D transposed convolution over an input with several channels + + Args: + input (array): Input array of shape ``(N, L, C_in)``. + weight (array): Weight array of shape ``(C_out, K, C_in)``. + stride (int, optional): Kernel stride. Default: ``1``. + padding (int, optional): Input padding. Default: ``0``. + dilation (int, optional): Kernel dilation. Default: ``1``. + output_padding (int, optional): Output padding. Default: ``0``. + groups (int, optional): Input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def conv_transpose2d( + input: array, + weight: array, + /, + stride: int | tuple[int, int] = ..., + padding: int | tuple[int, int] = ..., + dilation: int | tuple[int, int] = ..., + output_padding: int | tuple[int, int] = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 2D transposed convolution over an input with several channels + + Note: Only the default ``groups=1`` is currently supported. + + Args: + input (array): Input array of shape ``(N, H, W, C_in)``. + weight (array): Weight array of shape ``(C_out, KH, KW, C_in)``. + stride (int or tuple(int), optional): :obj:`tuple` of size 2 with + kernel strides. All spatial dimensions get the same stride if + only one number is specified. Default: ``1``. + padding (int or tuple(int), optional): :obj:`tuple` of size 2 with + symmetric input padding. All spatial dimensions get the same + padding if only one number is specified. Default: ``0``. + dilation (int or tuple(int), optional): :obj:`tuple` of size 2 with + kernel dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + output_padding (int or tuple(int), optional): :obj:`tuple` of size 2 with + output padding. All spatial dimensions get the same output + padding if only one number is specified. Default: ``0``. + groups (int, optional): input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def conv_transpose3d( + input: array, + weight: array, + /, + stride: int | tuple[int, int, int] = ..., + padding: int | tuple[int, int, int] = ..., + dilation: int | tuple[int, int, int] = ..., + output_padding: int | tuple[int, int, int] = ..., + groups: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + 3D transposed convolution over an input with several channels + + Note: Only the default ``groups=1`` is currently supported. + + Args: + input (array): Input array of shape ``(N, D, H, W, C_in)``. + weight (array): Weight array of shape ``(C_out, KD, KH, KW, C_in)``. + stride (int or tuple(int), optional): :obj:`tuple` of size 3 with + kernel strides. All spatial dimensions get the same stride if + only one number is specified. Default: ``1``. + padding (int or tuple(int), optional): :obj:`tuple` of size 3 with + symmetric input padding. All spatial dimensions get the same + padding if only one number is specified. Default: ``0``. + dilation (int or tuple(int), optional): :obj:`tuple` of size 3 with + kernel dilation. All spatial dimensions get the same dilation + if only one number is specified. Default: ``1`` + output_padding (int or tuple(int), optional): :obj:`tuple` of size 3 with + output padding. All spatial dimensions get the same output + padding if only one number is specified. Default: ``0``. + groups (int, optional): input feature groups. Default: ``1``. + + Returns: + array: The convolved array. + """ + +def convolve( + a: array, v: array, /, mode: str = ..., *, stream: Stream | Device | None = ... +) -> array: + """ + The discrete convolution of 1D arrays. + + If ``v`` is longer than ``a``, then they are swapped. + The conv filter is flipped following signal processing convention. + + Args: + a (array): 1D Input array. + v (array): 1D Input array. + mode (str, optional): {'full', 'valid', 'same'} + + Returns: + array: The convolved array. + """ + +def cos(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise cosine. + + Args: + a (array): Input array. + + Returns: + array: The cosine of ``a``. + """ + +def cosh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise hyperbolic cosine. + + Args: + a (array): Input array. + + Returns: + array: The hyperbolic cosine of ``a``. + """ + +cpu: DeviceType = ... + +def cummax( + a: array, + /, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return the cumulative maximum of the elements along the given axis. + + Args: + a (array): Input array + axis (int, optional): Optional axis to compute the cumulative maximum + over. If unspecified the cumulative maximum of the flattened array is + returned. + reverse (bool): Perform the cumulative maximum in reverse. + inclusive (bool): The i-th element of the output includes the i-th + element of the input. + + Returns: + array: The output array. + """ + +def cummin( + a: array, + /, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return the cumulative minimum of the elements along the given axis. + + Args: + a (array): Input array + axis (int, optional): Optional axis to compute the cumulative minimum + over. If unspecified the cumulative minimum of the flattened array is + returned. + reverse (bool): Perform the cumulative minimum in reverse. + inclusive (bool): The i-th element of the output includes the i-th + element of the input. + + Returns: + array: The output array. + """ + +def cumprod( + a: array, + /, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return the cumulative product of the elements along the given axis. + + Args: + a (array): Input array + axis (int, optional): Optional axis to compute the cumulative product + over. If unspecified the cumulative product of the flattened array is + returned. + reverse (bool): Perform the cumulative product in reverse. + inclusive (bool): The i-th element of the output includes the i-th + element of the input. + + Returns: + array: The output array. + """ + +def cumsum( + a: array, + /, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return the cumulative sum of the elements along the given axis. + + Args: + a (array): Input array + axis (int, optional): Optional axis to compute the cumulative sum + over. If unspecified the cumulative sum of the flattened array is + returned. + reverse (bool): Perform the cumulative sum in reverse. + inclusive (bool): The i-th element of the output includes the i-th + element of the input. + + Returns: + array: The output array. + """ + +class custom_function: + """ + Set up a function for custom gradient and vmap definitions. + + This class is meant to be used as a function decorator. Instances are + callables that behave identically to the wrapped function. However, when + a function transformation is used (e.g. computing gradients using + :func:`value_and_grad`) then the functions defined via + :meth:`custom_function.vjp`, :meth:`custom_function.jvp` and + :meth:`custom_function.vmap` are used instead of the default transformation. + + Note, all custom transformations are optional. Undefined transformations + fall back to the default behaviour. + + Example: + + .. code-block:: python + + import mlx.core as mx + + @mx.custom_function + def f(x, y): + return mx.sin(x) * y + + @f.vjp + def f_vjp(primals, cotangent, output): + x, y = primals + return cotan * mx.cos(x) * y, cotan * mx.sin(x) + + @f.jvp + def f_jvp(primals, tangents): + x, y = primals + dx, dy = tangents + return dx * mx.cos(x) * y + dy * mx.sin(x) + + @f.vmap + def f_vmap(inputs, axes): + x, y = inputs + ax, ay = axes + if ay != ax and ax is not None: + y = y.swapaxes(ay, ax) + return mx.sin(x) * y, (ax or ay) + + All ``custom_function`` instances behave as pure functions. Namely, any + variables captured will be treated as constants and no gradients will be + computed with respect to the captured arrays. For instance: + + .. code-block:: python + + import mlx.core as mx + + def g(x, y): + @mx.custom_function + def f(x): + return x * y + + @f.vjp + def f_vjp(x, dx, fx): + # Note that we have only x, dx and fx and nothing with respect to y + raise ValueError("Abort!") + + return f(x) + + x = mx.array(2.0) + y = mx.array(3.0) + print(g(x, y)) # prints 6.0 + print(mx.grad(g)(x, y)) # Raises exception + print(mx.grad(g, argnums=1)(x, y)) # prints 0.0 + """ + def __init__(self, f: Callable) -> None: ... + def __call__(self, *args, **kwargs) -> object: ... + def vjp(self, f: Callable): + """ + Define a custom vjp for the wrapped function. + + The vjp function takes three arguments: + + - *primals*: A pytree that contains all the positional arguments to + the function. It could be a single array, a tuple of arrays or a + full blown tuple of dicts of arrays etc. + - *cotangents*: A pytree that matches the structure of the output + but contains the cotangents (usually the gradients of the loss + function with respect to the outputs). + - *outputs*: The outputs of the function to be used to avoid + recomputing them for the gradient computation. + + The vjp function should return the same pytree structure as the + primals but containing the corresponding computed cotangents. + """ + + def jvp(self, f: Callable): + """ + Define a custom jvp for the wrapped function. + + The jvp function takes two arguments: + + - *primals*: A pytree that contains all the positional arguments to + the function. It could be a single array, a tuple of arrays or a + full blown tuple of dicts of arrays etc. + - *tangents*: A pytree that matches the structure of the inputs but + instead contains the gradients wrt to each input. Tangents could + be ``None`` if some inputs don't have an associated gradient. + + The jvp function should return the same pytree structure as the + outputs of the function but containing the tangents. + """ + + def vmap(self, f: Callable): + """ + Define a custom vectorization transformation for the wrapped function. + + The vmap function takes two arguments: + + - *inputs*: A pytree that contains all the positional arguments to + the function. It could be a single array, a tuple of arrays or a + full blown tuple of dicts of arrays etc. + - *axes*: A pytree that matches the structure of the inputs but + instead contains the vectorization axis for each input or + ``None`` if an input is not vectorized. + + The vmap function should return the outputs of the original + function but vectorized over the provided axes. It should also + return a pytree with the vectorization axes of each output. If some + outputs are no longer vectorized, then their vectorization axis + should be ``None``. + """ + +def default_device() -> Device: + """Get the default device.""" + +def default_stream(device: Device) -> Stream: + """Get the device's default stream.""" + +def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Convert angles from radians to degrees. + + Args: + a (array): Input array. + + Returns: + array: The angles in degrees. + """ + +def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]): + """ + Insert dependencies between arrays in the graph. The outputs are + identical to ``inputs`` but with dependencies on ``dependencies``. + + Args: + inputs (array or Sequence[array]): The input array or arrays. + dependencies (array or Sequence[array]): The array or arrays + to insert dependencies on. + + Returns: + array or Sequence[array]: The outputs which depend on dependencies. + """ + +def dequantize( + w: array, + /, + scales: array, + biases: array | None = ..., + group_size: int = ..., + bits: int = ..., + mode: str = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + r""" + Dequantize the matrix ``w`` using quantization parameters. + + Args: + w (array): Matrix to be dequantized + scales (array): The scales to use per ``group_size`` elements of ``w``. + biases (array, optional): The biases to use per ``group_size`` + elements of ``w``. Default: ``None``. + group_size (int, optional): The size of the group in ``w`` that shares a + scale and bias. Default: ``64``. + bits (int, optional): The number of bits occupied by each element in + ``w``. Default: ``4``. + mode (str, optional): The quantization mode. Default: ``"affine"``. + + Returns: + array: The dequantized version of ``w`` + + Notes: + The currently supported quantization modes are ``"affine"`` and ``mxfp4``. + + For ``affine`` quantization, given the notation in :func:`quantize`, + we compute :math:`w_i` from :math:`\hat{w_i}` and corresponding :math:`s` + and :math:`\beta` as follows + + .. math:: + + w_i = s \hat{w_i} + \beta + """ + +def diag(a: array, /, k: int = ..., *, stream: Stream | Device | None = ...) -> array: + """ + Extract a diagonal or construct a diagonal matrix. + If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the + :math:`k`-th diagonal. If ``a`` is 2-D then the :math:`k`-th diagonal is + returned. + + Args: + a (array): 1-D or 2-D input array. + k (int, optional): The diagonal to extract or construct. + Default: ``0``. + + Returns: + array: The extracted diagonal or the constructed diagonal matrix. + """ + +def diagonal( + a: array, + offset: int = ..., + axis1: int = ..., + axis2: int = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return specified diagonals. + + If ``a`` is 2-D, then a 1-D array containing the diagonal at the given + ``offset`` is returned. + + If ``a`` has more than two dimensions, then ``axis1`` and ``axis2`` + determine the 2D subarrays from which diagonals are extracted. The new + shape is the original shape with ``axis1`` and ``axis2`` removed and a + new dimension inserted at the end corresponding to the diagonal. + + Args: + a (array): Input array + offset (int, optional): Offset of the diagonal from the main diagonal. + Can be positive or negative. Default: ``0``. + axis1 (int, optional): The first axis of the 2-D sub-arrays from which + the diagonals should be taken. Default: ``0``. + axis2 (int, optional): The second axis of the 2-D sub-arrays from which + the diagonals should be taken. Default: ``1``. + + Returns: + array: The diagonals of the array. + """ + +def disable_compile() -> None: + """ + Globally disable compilation. Setting the environment variable + ``MLX_DISABLE_COMPILE`` can also be used to disable compilation. + """ + +def divide( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise division. + + Divide two arrays with numpy-style broadcasting semantics. Either or both + input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The quotient ``a / b``. + """ + +def divmod( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise quotient and remainder. + + The fuction ``divmod(a, b)`` is equivalent to but faster than + ``(a // b, a % b)``. The function uses numpy-style broadcasting + semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + tuple(array, array): The quotient ``a // b`` and remainder ``a % b``. + """ + +e: float = ... + +def einsum(subscripts: str, *operands, stream: Stream | Device | None = ...) -> array: + """ + Perform the Einstein summation convention on the operands. + + Args: + subscripts (str): The Einstein summation convention equation. + *operands (array): The input arrays. + + Returns: + array: The output array. + """ + +def einsum_path(subscripts: str, *operands): + """ + Compute the contraction order for the given Einstein summation. + + Args: + subscripts (str): The Einstein summation convention equation. + *operands (array): The input arrays. + + Returns: + tuple(list(tuple(int, int)), str): + The einsum path and a string containing information about the + chosen path. + """ + +def enable_compile() -> None: + """ + Globally enable compilation. This will override the environment + variable ``MLX_DISABLE_COMPILE`` if set. + """ + +def equal( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise equality. + + Equality comparison on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a == b``. + """ + +def erf(a: array, /, *, stream: Stream | Device | None = ...) -> array: + r""" + Element-wise error function. + + .. math:: + \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} \, dt + + Args: + a (array): Input array. + + Returns: + array: The error function of ``a``. + """ + +def erfinv(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise inverse of :func:`erf`. + + Args: + a (array): Input array. + + Returns: + array: The inverse error function of ``a``. + """ + +euler_gamma: float = ... + +type MX_ARRAY_TREE = ( + array + | list[MX_ARRAY_TREE] + | tuple[MX_ARRAY_TREE, ...] + | Mapping[Any, MX_ARRAY_TREE] +) + +def eval(*args: MX_ARRAY_TREE) -> None: + """ + Evaluate an :class:`array` or tree of :class:`array`. + + Args: + *args (arrays or trees of arrays): Each argument can be a single array + or a tree of arrays. If a tree is given the nodes can be a Python + :class:`list`, :class:`tuple` or :class:`dict`. Leaves which are not + arrays are ignored. + """ + +def exp(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise exponential. + + Args: + a (array): Input array. + + Returns: + array: The exponential of ``a``. + """ + +def expand_dims( + a: array, + /, + axis: int | Sequence[int], + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Add a size one dimension at the given axis. + + Args: + a (array): Input array. + axes (int or tuple(int)): The index of the inserted dimensions. + + Returns: + array: The array with inserted dimensions. + """ + +def expm1(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise exponential minus 1. + + Computes ``exp(x) - 1`` with greater precision for small ``x``. + + Args: + a (array): Input array. + + Returns: + array: The expm1 of ``a``. + """ + +def export_function( + arg0: object, fun: Callable, *args, shapeless: bool = ..., **kwargs +) -> None: + """ + Export an MLX function. + + Example input arrays must be provided to export a function. The example + inputs can be variable ``*args`` and ``**kwargs`` or a tuple of arrays + and/or dictionary of string keys with array values. + + .. warning:: + + This is part of an experimental API which is likely to + change in future versions of MLX. Functions exported with older + versions of MLX may not be compatible with future versions. + + Args: + file (str or Callable): Either a file path to export the function + to or a callback. + fun (Callable): A function which takes as input zero or more + :class:`array` and returns one or more :class:`array`. + *args (array): Example array inputs to the function. + shapeless (bool, optional): Whether or not the function allows + inputs with variable shapes. Default: ``False``. + **kwargs (array): Additional example keyword array inputs to the + function. + + Example: + + .. code-block:: python + + def fun(x, y): + return x + y + + x = mx.array(1) + y = mx.array([1, 2, 3]) + mx.export_function("fun.mlxfn", fun, x, y=y) + """ + +def export_to_dot(file: object, *args, **kwargs) -> None: + """ + Export a graph to DOT format for visualization. + + A variable number of output arrays can be provided for exporting + The graph exported will recursively include all unevaluated inputs of + the provided outputs. + + Args: + file (str): The file path to export to. + *args (array): The output arrays. + **kwargs (dict[str, array]): Provide some names for arrays in the + graph to make the result easier to parse. + + Example: + >>> a = mx.array(1) + mx.array(2) + >>> mx.export_to_dot("graph.dot", a) + >>> x = mx.array(1) + >>> y = mx.array(2) + >>> mx.export_to_dot("graph.dot", x + y, x=x, y=y) + """ + +def exporter(file: str, fun: Callable, *, shapeless: bool = ...) -> FunctionExporter: + """ + Make a callable object to export multiple traces of a function to a file. + + .. warning:: + + This is part of an experimental API which is likely to + change in future versions of MLX. Functions exported with older + versions of MLX may not be compatible with future versions. + + Args: + file (str): File path to export the function to. + shapeless (bool, optional): Whether or not the function allows + inputs with variable shapes. Default: ``False``. + + Example: + + .. code-block:: python + + def fun(*args): + return sum(args) + + with mx.exporter("fun.mlxfn", fun) as exporter: + exporter(mx.array(1)) + exporter(mx.array(1), mx.array(2)) + exporter(mx.array(1), mx.array(2), mx.array(3)) + """ + +def eye( + n: int, + m: int | None = ..., + k: int = ..., + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Create an identity matrix or a general diagonal matrix. + + Args: + n (int): The number of rows in the output. + m (int, optional): The number of columns in the output. Defaults to n. + k (int, optional): Index of the diagonal. Defaults to 0 (main diagonal). + dtype (Dtype, optional): Data type of the output array. Defaults to float32. + stream (Stream, optional): Stream or device. Defaults to None. + + Returns: + array: An array where all elements are equal to zero, except for the k-th diagonal, whose values are equal to one. + """ + +class finfo: + """Get information on floating-point types.""" + def __init__(self, arg: Dtype, /) -> None: ... + @property + def min(self) -> float: + """The smallest representable number.""" + + @property + def max(self) -> float: + """The largest representable number.""" + + @property + def eps(self) -> float: + """ + The difference between 1.0 and the next smallest + representable number larger than 1.0. + """ + + @property + def dtype(self) -> Dtype: + """The :obj:`Dtype`.""" + + def __repr__(self) -> str: ... + +def flatten( + a: array, + /, + start_axis: int = ..., + end_axis: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Flatten an array. + + The axes flattened will be between ``start_axis`` and ``end_axis``, + inclusive. Negative axes are supported. After converting negative axis to + positive, axes outside the valid range will be clamped to a valid value, + ``start_axis`` to ``0`` and ``end_axis`` to ``ndim - 1``. + + Args: + a (array): Input array. + start_axis (int, optional): The first dimension to flatten. Defaults to ``0``. + end_axis (int, optional): The last dimension to flatten. Defaults to ``-1``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The flattened array. + + Example: + >>> a = mx.array([[1, 2], [3, 4]]) + >>> mx.flatten(a) + array([1, 2, 3, 4], dtype=int32) + >>> + >>> mx.flatten(a, start_axis=0, end_axis=-1) + array([1, 2, 3, 4], dtype=int32) + """ + +float16: Dtype = ... +float32: Dtype = ... +float64: Dtype = ... +floating: DtypeCategory = ... + +def floor(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise floor. + + Args: + a (array): Input array. + + Returns: + array: The floor of ``a``. + """ + +def floor_divide( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise integer division. + + If either array is a floating point type then it is equivalent to + calling :func:`floor` after :func:`divide`. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The quotient ``a // b``. + """ + +def full( + shape: int | Sequence[int], + vals: scalar | array, + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Construct an array with the given value. + + Constructs an array of size ``shape`` filled with ``vals``. If ``vals`` + is an :obj:`array` it must be broadcastable to the given ``shape``. + + Args: + shape (int or list(int)): The shape of the output array. + vals (float or int or array): Values to fill the array with. + dtype (Dtype, optional): Data type of the output array. If + unspecified the output type is inferred from ``vals``. + + Returns: + array: The output array with the specified shape and values. + """ + +def gather_mm( + a: array, + b: array, + /, + lhs_indices: array, + rhs_indices: array, + *, + sorted_indices: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Matrix multiplication with matrix-level gather. + + Performs a gather of the operands with the given indices followed by a + (possibly batched) matrix multiplication of two arrays. This operation + is more efficient than explicitly applying a :func:`take` followed by a + :func:`matmul`. + + The indices ``lhs_indices`` and ``rhs_indices`` contain flat indices + along the batch dimensions (i.e. all but the last two dimensions) of + ``a`` and ``b`` respectively. + + For ``a`` with shape ``(A1, A2, ..., AS, M, K)``, ``lhs_indices`` + contains indices from the range ``[0, A1 * A2 * ... * AS)`` + + For ``b`` with shape ``(B1, B2, ..., BS, M, K)``, ``rhs_indices`` + contains indices from the range ``[0, B1 * B2 * ... * BS)`` + + If only one index is passed and it is sorted, the ``sorted_indices`` + flag can be passed for a possible faster implementation. + + Args: + a (array): Input array. + b (array): Input array. + lhs_indices (array, optional): Integer indices for ``a``. Default: ``None`` + rhs_indices (array, optional): Integer indices for ``b``. Default: ``None`` + sorted_indices (bool, optional): May allow a faster implementation + if the passed indices are sorted. Default: ``False``. + + Returns: + array: The output array. + """ + +def gather_qmm( + x: array, + w: array, + /, + scales: array, + biases: array | None = ..., + lhs_indices: array | None = ..., + rhs_indices: array | None = ..., + transpose: bool = ..., + group_size: int = ..., + bits: int = ..., + mode: str = ..., + *, + sorted_indices: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Perform quantized matrix multiplication with matrix-level gather. + + This operation is the quantized equivalent to :func:`gather_mm`. + Similar to :func:`gather_mm`, the indices ``lhs_indices`` and + ``rhs_indices`` contain flat indices along the batch dimensions (i.e. + all but the last two dimensions) of ``x`` and ``w`` respectively. + + Note that ``scales`` and ``biases`` must have the same batch dimensions + as ``w`` since they represent the same quantized matrix. + + Args: + x (array): Input array + w (array): Quantized matrix packed in unsigned integers + scales (array): The scales to use per ``group_size`` elements of ``w`` + biases (array, optional): The biases to use per ``group_size`` + elements of ``w``. Default: ``None``. + lhs_indices (array, optional): Integer indices for ``x``. Default: ``None``. + rhs_indices (array, optional): Integer indices for ``w``. Default: ``None``. + transpose (bool, optional): Defines whether to multiply with the + transposed ``w`` or not, namely whether we are performing + ``x @ w.T`` or ``x @ w``. Default: ``True``. + group_size (int, optional): The size of the group in ``w`` that + shares a scale and bias. Default: ``64``. + bits (int, optional): The number of bits occupied by each element in + ``w``. Default: ``4``. + mode (str, optional): The quantization mode. Default: ``"affine"``. + sorted_indices (bool, optional): May allow a faster implementation + if the passed indices are sorted. Default: ``False``. + + Returns: + array: The result of the multiplication of ``x`` with ``w`` + after gathering using ``lhs_indices`` and ``rhs_indices``. + """ + +generic: DtypeCategory = ... + +def get_active_memory() -> int: + """ + Get the actively used memory in bytes. + + Note, this will not always match memory use reported by the system because + it does not include cached memory buffers. + """ + +def get_cache_memory() -> int: + """ + Get the cache size in bytes. + + The cache includes memory not currently used that has not been returned + to the system allocator. + """ + +def get_peak_memory() -> int: + """ + Get the peak amount of used memory in bytes. + + The maximum memory used recorded from the beginning of the program + execution or since the last call to :func:`reset_peak_memory`. + """ + +gpu: DeviceType = ... + +def grad( + fun: Callable, + argnums: int | Sequence[int] | None = ..., + argnames: str | Sequence[str] = ..., +) -> Callable: + """ + Returns a function which computes the gradient of ``fun``. + + Args: + fun (Callable): A function which takes a variable number of + :class:`array` or trees of :class:`array` and returns + a scalar output :class:`array`. + argnums (int or list(int), optional): Specify the index (or indices) + of the positional arguments of ``fun`` to compute the gradient + with respect to. If neither ``argnums`` nor ``argnames`` are + provided ``argnums`` defaults to ``0`` indicating ``fun``'s first + argument. + argnames (str or list(str), optional): Specify keyword arguments of + ``fun`` to compute gradients with respect to. It defaults to [] so + no gradients for keyword arguments by default. + + Returns: + Callable: A function which has the same input arguments as ``fun`` and + returns the gradient(s). + """ + +def greater( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise greater than. + + Strict greater than on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a > b``. + """ + +def greater_equal( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise greater or equal. + + Greater than or equal on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a >= b``. + """ + +def hadamard_transform( + a: array, scale: float | None = ..., stream: Stream | Device | None = ... +) -> array: + """ + Perform the Walsh-Hadamard transform along the final axis. + + Equivalent to: + + .. code-block:: python + + from scipy.linalg import hadamard + + y = (hadamard(len(x)) @ x) * scale + + Supports sizes ``n = m*2^k`` for ``m`` in ``(1, 12, 20, 28)`` and ``2^k + <= 8192`` for float32 and ``2^k <= 16384`` for float16/bfloat16. + + Args: + a (array): Input array or scalar. + scale (float): Scale the output by this factor. + Defaults to ``1/sqrt(a.shape[-1])`` so that the Hadamard matrix is orthonormal. + + Returns: + array: The transformed array. + """ + +def identity( + n: int, dtype: Dtype | None = ..., *, stream: Stream | Device | None = ... +) -> array: + """ + Create a square identity matrix. + + Args: + n (int): The number of rows and columns in the output. + dtype (Dtype, optional): Data type of the output array. Defaults to float32. + stream (Stream, optional): Stream or device. Defaults to None. + + Returns: + array: An identity matrix of size n x n. + """ + +class iinfo: + """Get information on integer types.""" + def __init__(self, arg: Dtype, /) -> None: ... + @property + def min(self) -> int: + """The smallest representable number.""" + + @property + def max(self) -> int: + """The largest representable number.""" + + @property + def dtype(self) -> Dtype: + """The :obj:`Dtype`.""" + + def __repr__(self) -> str: ... + +def imag(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Returns the imaginary part of a complex array. + + Args: + a (array): Input array. + + Returns: + array: The imaginary part of ``a``. + """ + +def import_function(file: str) -> Callable: + """ + Import a function from a file. + + The imported function can be called either with ``*args`` and + ``**kwargs`` or with a tuple of arrays and/or dictionary of string + keys with array values. Imported functions always return a tuple of + arrays. + + .. warning:: + + This is part of an experimental API which is likely to + change in future versions of MLX. Functions exported with older + versions of MLX may not be compatible with future versions. + + Args: + file (str): The file path to import the function from. + + Returns: + Callable: The imported function. + + Example: + >>> fn = mx.import_function("function.mlxfn") + >>> out = fn(a, b, x=x, y=y)[0] + >>> + >>> out = fn((a, b), {"x": x, "y": y}[0] + """ + +inexact: DtypeCategory = ... +inf: float = ... + +def inner(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes. + + Args: + a (array): Input array + b (array): Input array + + Returns: + array: The inner product. + """ + +int16: Dtype = ... +int32: Dtype = ... +int64: Dtype = ... +int8: Dtype = ... +integer: DtypeCategory = ... + +def is_available(device: Device) -> bool: + """Check if a back-end is available for the given device.""" + +def isclose( + a: array, + b: array, + /, + rtol: float = ..., + atol: float = ..., + *, + equal_nan: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Returns a boolean array where two arrays are element-wise equal within a tolerance. + + Infinite values are considered equal if they have the same sign, NaN values are + not equal unless ``equal_nan`` is ``True``. + + Two values are considered equal if: + + .. code-block:: + + abs(a - b) <= (atol + rtol * abs(b)) + + Note unlike :func:`array_equal`, this function supports numpy-style + broadcasting. + + Args: + a (array): Input array. + b (array): Input array. + rtol (float): Relative tolerance. + atol (float): Absolute tolerance. + equal_nan (bool): If ``True``, NaNs are considered equal. + Defaults to ``False``. + + Returns: + array: The boolean output scalar indicating if the arrays are close. + """ + +def isfinite(a: array, stream: Stream | Device | None = ...) -> array: + """ + Return a boolean array indicating which elements are finite. + + An element is finite if it is not infinite or NaN. + + Args: + a (array): Input array. + + Returns: + array: The boolean array indicating which elements are finite. + """ + +def isinf(a: array, stream: Stream | Device | None = ...) -> array: + """ + Return a boolean array indicating which elements are +/- inifnity. + + Args: + a (array): Input array. + + Returns: + array: The boolean array indicating which elements are +/- infinity. + """ + +def isnan(a: array, stream: Stream | Device | None = ...) -> array: + """ + Return a boolean array indicating which elements are NaN. + + Args: + a (array): Input array. + + Returns: + array: The boolean array indicating which elements are NaN. + """ + +def isneginf(a: array, stream: Stream | Device | None = ...) -> array: + """ + Return a boolean array indicating which elements are negative infinity. + + Args: + a (array): Input array. + stream (Stream | Device | None): Optional stream or device. + + Returns: + array: The boolean array indicating which elements are negative infinity. + """ + +def isposinf(a: array, stream: Stream | Device | None = ...) -> array: + """ + Return a boolean array indicating which elements are positive infinity. + + Args: + a (array): Input array. + stream (Stream | Device | None): Optional stream or device. + + Returns: + array: The boolean array indicating which elements are positive infinity. + """ + +def issubdtype(arg1: Dtype | DtypeCategory, arg2: Dtype | DtypeCategory) -> bool: + """ + Check if a :obj:`Dtype` or :obj:`DtypeCategory` is a subtype + of another. + + Args: + arg1 (Dtype | DtypeCategory: First dtype or category. + arg2 (Dtype | DtypeCategory: Second dtype or category. + + Returns: + bool: + A boolean indicating if the first input is a subtype of the + second input. + + Example: + + >>> ints = mx.array([1, 2, 3], dtype=mx.int32) + >>> mx.issubdtype(ints.dtype, mx.integer) + True + >>> mx.issubdtype(ints.dtype, mx.floating) + False + + >>> floats = mx.array([1, 2, 3], dtype=mx.float32) + >>> mx.issubdtype(floats.dtype, mx.integer) + False + >>> mx.issubdtype(floats.dtype, mx.floating) + True + + Similar types of different sizes are not subdtypes of each other: + + >>> mx.issubdtype(mx.float64, mx.float32) + False + >>> mx.issubdtype(mx.float32, mx.float64) + False + + but both are subtypes of `floating`: + + >>> mx.issubdtype(mx.float64, mx.floating) + True + >>> mx.issubdtype(mx.float32, mx.floating) + True + + For convenience, dtype-like objects are allowed too: + + >>> mx.issubdtype(mx.float32, mx.inexact) + True + >>> mx.issubdtype(mx.signedinteger, mx.floating) + False + """ + +def jvp( + fun: Callable, primals: list[array], tangents: list[array] +) -> tuple[list[array], list[array]]: + """ + Compute the Jacobian-vector product. + + This computes the product of the Jacobian of a function ``fun`` evaluated + at ``primals`` with the ``tangents``. + + Args: + fun (Callable): A function which takes a variable number of :class:`array` + and returns a single :class:`array` or list of :class:`array`. + primals (list(array)): A list of :class:`array` at which to + evaluate the Jacobian. + tangents (list(array)): A list of :class:`array` which are the + "vector" in the Jacobian-vector product. The ``tangents`` should be the + same in number, shape, and type as the inputs of ``fun`` (i.e. the ``primals``). + + Returns: + list(array): A list of the Jacobian-vector products which + is the same in number, shape, and type of the inputs to ``fun``. + """ + +def kron(a: array, b: array, *, stream: Stream | Device | None = ...) -> array: + """ + Compute the Kronecker product of two arrays ``a`` and ``b``. + + Args: + a (array): The first input array. + b (array): The second input array. + stream (Stream | Device | None, optional): Optional stream or + device for execution. Default: ``None``. + + Returns: + array: The Kronecker product of ``a`` and ``b``. + + Examples: + >>> a = mx.array([[1, 2], [3, 4]]) + >>> b = mx.array([[0, 5], [6, 7]]) + >>> result = mx.kron(a, b) + >>> print(result) + array([[0, 5, 0, 10], + [6, 7, 12, 14], + [0, 15, 0, 20], + [18, 21, 24, 28]], dtype=int32) + """ + +def left_shift( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise left shift. + + Shift the bits of the first input to the left by the second using + numpy-style broadcasting semantics. Either or both input arrays can + also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The bitwise left shift ``a << b``. + """ + +def less( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise less than. + + Strict less than on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a < b``. + """ + +def less_equal( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise less than or equal. + + Less than or equal on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a <= b``. + """ + +def linspace( + start, + stop, + num: int | None = ..., + dtype: Dtype | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. + + Args: + start (scalar): Starting value. + stop (scalar): Stopping value. + num (int, optional): Number of samples, defaults to ``50``. + dtype (Dtype, optional): Specifies the data type of the output, + default to ``float32``. + + Returns: + array: The range of values. + """ + +def load( + file: str | pathlib.Path, + /, + format: str | None = ..., + return_metadata: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array | dict[str, array]: + """ + Load array(s) from a binary file. + + The supported formats are ``.npy``, ``.npz``, ``.safetensors``, and + ``.gguf``. + + Args: + file (str, pathlib.Path): File in which the array is saved. + format (str, optional): Format of the file. If ``None``, the + format is inferred from the file extension. Supported formats: + ``npy``, ``npz``, and ``safetensors``. Default: ``None``. + return_metadata (bool, optional): Load the metadata for formats + which support matadata. The metadata will be returned as an + additional dictionary. Default: ``False``. + Returns: + array or dict: + A single array if loading from a ``.npy`` file or a dict + mapping names to arrays if loading from a ``.npz`` or + ``.safetensors`` file. If ``return_metadata`` is ``True`` an + additional dictionary of metadata will be returned. + + Warning: + + When loading unsupported quantization formats from GGUF, tensors + will automatically cast to ``mx.float16`` + """ + +def log(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise natural logarithm. + + Args: + a (array): Input array. + + Returns: + array: The natural logarithm of ``a``. + """ + +def log10(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise base-10 logarithm. + + Args: + a (array): Input array. + + Returns: + array: The base-10 logarithm of ``a``. + """ + +def log1p(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise natural log of one plus the array. + + Args: + a (array): Input array. + + Returns: + array: The natural logarithm of one plus ``a``. + """ + +def log2(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise base-2 logarithm. + + Args: + a (array): Input array. + + Returns: + array: The base-2 logarithm of ``a``. + """ + +def logaddexp( + a: scalar | array, + b: scalar | array, + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise log-add-exp. + + This is a numerically stable log-add-exp of two arrays with numpy-style + broadcasting semantics. Either or both input arrays can also be scalars. + + The computation is is a numerically stable version of ``log(exp(a) + exp(b))``. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The log-add-exp of ``a`` and ``b``. + """ + +def logcumsumexp( + a: array, + /, + axis: int | None = ..., + *, + reverse: bool = ..., + inclusive: bool = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Return the cumulative logsumexp of the elements along the given axis. + + Args: + a (array): Input array + axis (int, optional): Optional axis to compute the cumulative logsumexp + over. If unspecified the cumulative logsumexp of the flattened array is + returned. + reverse (bool): Perform the cumulative logsumexp in reverse. + inclusive (bool): The i-th element of the output includes the i-th + element of the input. + + Returns: + array: The output array. + """ + +def logical_and( + a: array, b: array, /, *, stream: Stream | Device | None = ... +) -> array: + """ + Element-wise logical and. + + Args: + a (array): First input array or scalar. + b (array): Second input array or scalar. + + Returns: + array: The boolean array containing the logical and of ``a`` and ``b``. + """ + +def logical_not(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise logical not. + + Args: + a (array): Input array or scalar. + + Returns: + array: The boolean array containing the logical not of ``a``. + """ + +def logical_or(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise logical or. + + Args: + a (array): First input array or scalar. + b (array): Second input array or scalar. + + Returns: + array: The boolean array containing the logical or of ``a`` and ``b``. + """ + +def logsumexp( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + A `log-sum-exp` reduction over the given axes. + + The log-sum-exp reduction is a numerically stable version of: + + .. code-block:: + + log(sum(exp(a), axis)) + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def matmul(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Matrix multiplication. + + Perform the (possibly batched) matrix multiplication of two arrays. This function supports + broadcasting for arrays with more than two dimensions. + + - If the first array is 1-D then a 1 is prepended to its shape to make it + a matrix. Similarly if the second array is 1-D then a 1 is appended to its + shape to make it a matrix. In either case the singleton dimension is removed + from the result. + - A batched matrix multiplication is performed if the arrays have more than + 2 dimensions. The matrix dimensions for the matrix product are the last + two dimensions of each input. + - All but the last two dimensions of each input are broadcast with one another using + standard numpy-style broadcasting semantics. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The matrix product of ``a`` and ``b``. + """ + +def max( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + A `max` reduction over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def maximum( + a: scalar | array, + b: scalar | array, + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise maximum. + + Take the element-wise max of two arrays with numpy-style broadcasting + semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The max of ``a`` and ``b``. + """ + +def mean( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Compute the mean(s) over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array of means. + """ + +def meshgrid( + *arrays: array, + sparse: bool | None = ..., + indexing: str | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate multidimensional coordinate grids from 1-D coordinate arrays + + Args: + *arrays (array): Input arrays. + sparse (bool, optional): If ``True``, a sparse grid is returned in which each output + array has a single non-zero element. If ``False``, a dense grid is returned. + Defaults to ``False``. + indexing (str, optional): Cartesian ('xy') or matrix ('ij') indexing of the output arrays. + Defaults to ``'xy'``. + + Returns: + list(array): The output arrays. + """ + +def min( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + A `min` reduction over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def minimum( + a: scalar | array, + b: scalar | array, + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise minimum. + + Take the element-wise min of two arrays with numpy-style broadcasting + semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The min of ``a`` and ``b``. + """ + +def moveaxis( + a: array, + /, + source: int, + destination: int, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Move an axis to a new position. + + Args: + a (array): Input array. + source (int): Specifies the source axis. + destination (int): Specifies the destination axis. + + Returns: + array: The array with the axis moved. + """ + +def multiply( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise multiplication. + + Multiply two arrays with numpy-style broadcasting semantics. Either or both + input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The multiplication ``a * b``. + """ + +nan: float = ... + +def nan_to_num( + a: scalar | array, + nan: float = ..., + posinf: float | None = ..., + neginf: float | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Replace NaN and Inf values with finite numbers. + + Args: + a (array): Input array + nan (float, optional): Value to replace NaN with. Default: ``0``. + posinf (float, optional): Value to replace positive infinities + with. If ``None``, defaults to largest finite value for the + given data type. Default: ``None``. + neginf (float, optional): Value to replace negative infinities + with. If ``None``, defaults to the negative of the largest + finite value for the given data type. Default: ``None``. + + Returns: + array: Output array with NaN and Inf replaced. + """ + +def negative(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise negation. + + Args: + a (array): Input array. + + Returns: + array: The negative of ``a``. + """ + +def new_stream(device: Device) -> Stream: + """Make a new stream on the given device.""" + +newaxis: None = ... + +def not_equal( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise not equal. + + Not equal comparison on two arrays with numpy-style broadcasting semantics. + Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The element-wise comparison ``a != b``. + """ + +number: DtypeCategory = ... + +def ones( + shape: int | Sequence[int], + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Construct an array of ones. + + Args: + shape (int or list(int)): The shape of the output array. + dtype (Dtype, optional): Data type of the output array. If + unspecified the output type defaults to ``float32``. + + Returns: + array: The array of ones with the specified shape. + """ + +def ones_like(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + An array of ones like the input. + + Args: + a (array): The input to take the shape and type from. + + Returns: + array: The output array filled with ones. + """ + +def outer(a: array, b: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand. + + Args: + a (array): Input array + b (array): Input array + + Returns: + array: The outer product. + """ + +def pad( + a: array, + pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], + mode: Literal["constant", "edge"] = ..., + constant_values: scalar | array = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Pad an array with a constant value + + Args: + a (array): Input array. + pad_width (int, tuple(int), tuple(int, int) or list(tuple(int, int))): Number of padded + values to add to the edges of each axis:``((before_1, after_1), + (before_2, after_2), ..., (before_N, after_N))``. If a single pair + of integers is passed then ``(before_i, after_i)`` are all the same. + If a single integer or tuple with a single integer is passed then + all axes are extended by the same number on each side. + mode: Padding mode. One of the following strings: + "constant" (default): Pads with a constant value. + "edge": Pads with the edge values of array. + constant_value (array or scalar, optional): Optional constant value + to pad the edges of the array with. + + Returns: + array: The padded array. + """ + +def partition( + a: array, + /, + kth: int, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Returns a partitioned copy of the array such that the smaller ``kth`` + elements are first. + + The ordering of the elements in partitions is undefined. + + Args: + a (array): Input array. + kth (int): Element at the ``kth`` index will be in its sorted + position in the output. All elements before the kth index will + be less or equal to the ``kth`` element and all elements after + will be greater or equal to the ``kth`` element in the output. + axis (int or None, optional): Optional axis to partition over. + If ``None``, this partitions over the flattened array. + If unspecified, it defaults to ``-1``. + + Returns: + array: The partitioned array. + """ + +def permute_dims( + a: array, + /, + axes: Sequence[int] | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """See :func:`transpose`.""" + +pi: float = ... + +def power( + a: scalar | array, + b: scalar | array, + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise power operation. + + Raise the elements of a to the powers in elements of b with numpy-style + broadcasting semantics. Either or both input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: Bases of ``a`` raised to powers in ``b``. + """ + +def prod( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + An product reduction over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def put_along_axis( + a: array, + /, + indices: array, + values: array, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Put values along an axis at the specified indices. + + Args: + a (array): Destination array. + indices (array): Indices array. These should be broadcastable with + the input array excluding the `axis` dimension. + values (array): Values array. These should be broadcastable with + the indices. + + axis (int or None): Axis in the destination to put the values to. If + ``axis == None`` the destination is flattened prior to the put + operation. + + Returns: + array: The output array. + """ + +def quantize( + w: array, + /, + group_size: int = ..., + bits: int = ..., + mode: str = ..., + *, + stream: Stream | Device | None = ..., +) -> tuple[array, array, array]: + r""" + Quantize the matrix ``w`` using ``bits`` bits per element. + + Note, every ``group_size`` elements in a row of ``w`` are quantized + together. Hence, number of columns of ``w`` should be divisible by + ``group_size``. In particular, the rows of ``w`` are divided into groups of + size ``group_size`` which are quantized together. + + .. warning:: + + ``quantize`` currently only supports 2D inputs with the second + dimension divisible by ``group_size`` + + The supported quantization modes are ``"affine"`` and ``"mxfp4"``. They + are described in more detail below. + + Args: + w (array): Matrix to be quantized + group_size (int, optional): The size of the group in ``w`` that shares a + scale and bias. Default: ``64``. + bits (int, optional): The number of bits occupied by each element of + ``w`` in the returned quantized matrix. Default: ``4``. + mode (str, optional): The quantization mode. Default: ``"affine"``. + + Returns: + tuple: A tuple with either two or three elements containing: + + * w_q (array): The quantized version of ``w`` + * scales (array): The quantization scales + * biases (array): The quantization biases (returned for ``mode=="affine"``). + + Notes: + The ``affine`` mode quantizes groups of :math:`g` consecutive + elements in a row of ``w``. For each group the quantized + representation of each element :math:`\hat{w_i}` is computed as follows: + + .. math:: + + \begin{aligned} + \alpha &= \max_i w_i \\ + \beta &= \min_i w_i \\ + s &= \frac{\alpha - \beta}{2^b - 1} \\ + \hat{w_i} &= \textrm{round}\left( \frac{w_i - \beta}{s}\right). + \end{aligned} + + After the above computation, :math:`\hat{w_i}` fits in :math:`b` bits + and is packed in an unsigned 32-bit integer from the lower to upper + bits. For instance, for 4-bit quantization we fit 8 elements in an + unsigned 32 bit integer where the 1st element occupies the 4 least + significant bits, the 2nd bits 4-7 etc. + + To dequantize the elements of ``w``, we also save :math:`s` and + :math:`\beta` which are the returned ``scales`` and + ``biases`` respectively. + + The ``mxfp4`` mode similarly quantizes groups of :math:`g` elements + of ``w``. For ``mxfp4`` the group size must be ``32``. The elements + are quantized to 4-bit precision floating-point values (E2M1) with a + shared 8-bit scale per group. Unlike ``affine`` quantization, + ``mxfp4`` does not have a bias value. More details on the format can + be found in the `specification `_. + """ + +def quantized_matmul( + x: array, + w: array, + /, + scales: array, + biases: array | None = ..., + transpose: bool = ..., + group_size: int = ..., + bits: int = ..., + mode: str = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Perform the matrix multiplication with the quantized matrix ``w``. The + quantization uses one floating point scale and bias per ``group_size`` of + elements. Each element in ``w`` takes ``bits`` bits and is packed in an + unsigned 32 bit integer. + + Args: + x (array): Input array + w (array): Quantized matrix packed in unsigned integers + scales (array): The scales to use per ``group_size`` elements of ``w`` + biases (array, optional): The biases to use per ``group_size`` + elements of ``w``. Default: ``None``. + transpose (bool, optional): Defines whether to multiply with the + transposed ``w`` or not, namely whether we are performing + ``x @ w.T`` or ``x @ w``. Default: ``True``. + group_size (int, optional): The size of the group in ``w`` that + shares a scale and bias. Default: ``64``. + bits (int, optional): The number of bits occupied by each element in + ``w``. Default: ``4``. + mode (str, optional): The quantization mode. Default: ``"affine"``. + + Returns: + array: The result of the multiplication of ``x`` with ``w``. + """ + +def radians(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Convert angles from degrees to radians. + + Args: + a (array): Input array. + + Returns: + array: The angles in radians. + """ + +def real(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Returns the real part of a complex array. + + Args: + a (array): Input array. + + Returns: + array: The real part of ``a``. + """ + +def reciprocal(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise reciprocal. + + Args: + a (array): Input array. + + Returns: + array: The reciprocal of ``a``. + """ + +def remainder( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise remainder of division. + + Computes the remainder of dividing a with b with numpy-style + broadcasting semantics. Either or both input arrays can also be + scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The remainder of ``a // b``. + """ + +def repeat( + array: array, + repeats: int, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Repeat an array along a specified axis. + + Args: + array (array): Input array. + repeats (int): The number of repetitions for each element. + axis (int, optional): The axis in which to repeat the array along. If + unspecified it uses the flattened array of the input and repeats + along axis 0. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: The resulting repeated array. + """ + +def reset_peak_memory() -> None: + """Reset the peak memory to zero.""" + +def reshape( + a: array, /, shape: Sequence[int], *, stream: Stream | Device | None = ... +) -> array: + """ + Reshape an array while preserving the size. + + Args: + a (array): Input array. + shape (tuple(int)): New shape. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The reshaped array. + """ + +def right_shift( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise right shift. + + Shift the bits of the first input to the right by the second using + numpy-style broadcasting semantics. Either or both input arrays can + also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The bitwise right shift ``a >> b``. + """ + +def roll( + a: array, + shift: int | tuple[int], + axis: int | tuple[int] | None = ..., + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Roll array elements along a given axis. + + Elements that are rolled beyond the end of the array are introduced at + the beggining and vice-versa. + + If the axis is not provided the array is flattened, rolled and then the + shape is restored. + + Args: + a (array): Input array + shift (int or tuple(int)): The number of places by which elements + are shifted. If positive the array is rolled to the right, if + negative it is rolled to the left. If an int is provided but the + axis is a tuple then the same value is used for all axes. + axis (int or tuple(int), optional): The axis or axes along which to + roll the elements. + """ + +def round( + a: array, /, decimals: int = ..., stream: Stream | Device | None = ... +) -> array: + """ + Round to the given number of decimals. + + Basically performs: + + .. code-block:: python + + s = 10**decimals + x = round(x * s) / s + + Args: + a (array): Input array + decimals (int): Number of decimal places to round to. (default: 0) + + Returns: + array: An array of the same type as ``a`` rounded to the + given number of decimals. + """ + +def rsqrt(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise reciprocal and square root. + + Args: + a (array): Input array. + + Returns: + array: One over the square root of ``a``. + """ + +def save(file: str | pathlib.Path, arr: array) -> None: + """ + Save the array to a binary file in ``.npy`` format. + + Args: + file (str, pathlib.Path): File to which the array is saved + arr (array): Array to be saved. + """ + +def save_gguf( + file: str | pathlib.Path, + arrays: dict[str, array], + metadata: dict[str, array | str | list[str]], +): + """ + Save array(s) to a binary file in ``.gguf`` format. + + See the `GGUF documentation + `_ for + more information on the format. + + Args: + file (file, str, pathlib.Path): File in which the array is saved. + arrays (dict(str, array)): The dictionary of names to arrays to + be saved. + metadata (dict(str, array | str | list(str))): The dictionary + of metadata to be saved. The values can be a scalar or 1D + obj:`array`, a :obj:`str`, or a :obj:`list` of :obj:`str`. + """ + +def save_safetensors( + file: str | pathlib.Path, + arrays: dict[str, array], + metadata: dict[str, str] | None = ..., +): + """ + Save array(s) to a binary file in ``.safetensors`` format. + + See the `Safetensors documentation + `_ for more + information on the format. + + Args: + file (file, str, pathlib.Path): File in which the array is saved. + arrays (dict(str, array)): The dictionary of names to arrays to + be saved. + metadata (dict(str, str), optional): The dictionary of + metadata to be saved. + """ + +def savez(file: str | pathlib.Path, *args, **kwargs): + """ + Save several arrays to a binary file in uncompressed ``.npz`` + format. + + .. code-block:: python + + import mlx.core as mx + + x = mx.ones((10, 10)) + mx.savez("my_path.npz", x=x) + + import mlx.nn as nn + from mlx.utils import tree_flatten + + model = nn.TransformerEncoder(6, 128, 4) + flat_params = tree_flatten(model.parameters()) + mx.savez("model.npz", **dict(flat_params)) + + Args: + file (file, str, pathlib.Path): Path to file to which the arrays are saved. + *args (arrays): Arrays to be saved. + **kwargs (arrays): Arrays to be saved. Each array will be saved + with the associated keyword as the output file name. + """ + +def savez_compressed(file: str | pathlib.Path, *args, **kwargs): + """ + Save several arrays to a binary file in compressed ``.npz`` format. + + Args: + file (file, str, pathlib.Path): Path to file to which the arrays are saved. + *args (arrays): Arrays to be saved. + **kwargs (arrays): Arrays to be saved. Each array will be saved + with the associated keyword as the output file name. + """ + +def segmented_mm( + a: array, b: array, /, segments: array, *, stream: Stream | Device | None = ... +) -> array: + """ + Perform a matrix multiplication but segment the inner dimension and + save the result for each segment separately. + + Args: + a (array): Input array of shape ``MxK``. + b (array): Input array of shape ``KxN``. + segments (array): The offsets into the inner dimension for each segment. + + Returns: + array: The result per segment of shape ``MxN``. + """ + +def set_cache_limit(limit: int) -> int: + """ + Set the free cache limit. + + If using more than the given limit, free memory will be reclaimed + from the cache on the next allocation. To disable the cache, set + the limit to ``0``. + + The cache limit defaults to the memory limit. See + :func:`set_memory_limit` for more details. + + Args: + limit (int): The cache limit in bytes. + + Returns: + int: The previous cache limit in bytes. + """ + +def set_default_device(device: Device | DeviceType) -> None: + """Set the default device.""" + +def set_default_stream(stream: Stream) -> None: + """ + Set the default stream. + + This will make the given stream the default for the + streams device. It will not change the default device. + + Args: + stream (stream): Stream to make the default. + """ + +def set_memory_limit(limit: int) -> int: + """ + Set the memory limit. + + The memory limit is a guideline for the maximum amount of memory to use + during graph evaluation. If the memory limit is exceeded and there is no + more RAM (including swap when available) allocations will result in an + exception. + + When metal is available the memory limit defaults to 1.5 times the + maximum recommended working set size reported by the device. + + Args: + limit (int): Memory limit in bytes. + + Returns: + int: The previous memory limit in bytes. + """ + +def set_wired_limit(limit: int) -> int: + """ + Set the wired size limit. + + .. note:: + * This function is only useful on macOS 15.0 or higher. + * The wired limit should remain strictly less than the total + memory size. + + The wired limit is the total size in bytes of memory that will be kept + resident. The default value is ``0``. + + Setting a wired limit larger than system wired limit is an error. You can + increase the system wired limit with: + + .. code-block:: + + sudo sysctl iogpu.wired_limit_mb= + + Use :func:`device_info` to query the system wired limit + (``"max_recommended_working_set_size"``) and the total memory size + (``"memory_size"``). + + Args: + limit (int): The wired limit in bytes. + + Returns: + int: The previous wired limit in bytes. + """ + +def sigmoid(a: array, /, *, stream: Stream | Device | None = ...) -> array: + r""" + Element-wise logistic sigmoid. + + The logistic sigmoid function is: + + .. math:: + \mathrm{sigmoid}(x) = \frac{1}{1 + e^{-x}} + + Args: + a (array): Input array. + + Returns: + array: The logistic sigmoid of ``a``. + """ + +def sign(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise sign. + + Args: + a (array): Input array. + + Returns: + array: The sign of ``a``. + """ + +signedinteger: DtypeCategory = ... + +def sin(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise sine. + + Args: + a (array): Input array. + + Returns: + array: The sine of ``a``. + """ + +def sinh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise hyperbolic sine. + + Args: + a (array): Input array. + + Returns: + array: The hyperbolic sine of ``a``. + """ + +def slice( + a: array, + start_indices: array, + axes: Sequence[int], + slice_size: Sequence[int], + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Extract a sub-array from the input array. + + Args: + a (array): Input array + start_indices (array): The index location to start the slice at. + axes (tuple(int)): The axes corresponding to the indices in ``start_indices``. + slice_size (tuple(int)): The size of the slice. + + Returns: + array: The sliced output array. + + Example: + + >>> a = mx.array([[1, 2, 3], [4, 5, 6]]) + >>> mx.slice(a, start_indices=mx.array(1), axes=(0,), slice_size=(1, 2)) + array([[4, 5]], dtype=int32) + >>> + >>> mx.slice(a, start_indices=mx.array(1), axes=(1,), slice_size=(2, 1)) + array([[2], + [5]], dtype=int32) + """ + +def slice_update( + a: array, + update: array, + start_indices: array, + axes: Sequence[int], + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Update a sub-array of the input array. + + Args: + a (array): The input array to update + update (array): The update array. + start_indices (array): The index location to start the slice at. + axes (tuple(int)): The axes corresponding to the indices in ``start_indices``. + + Returns: + array: The output array with the same shape and type as the input. + + Example: + + >>> a = mx.zeros((3, 3)) + >>> mx.slice_update(a, mx.ones((1, 2)), start_indices=mx.array(1, 1), axes=(0, 1)) + array([[0, 0, 0], + [0, 1, 0], + [0, 1, 0]], dtype=float32) + """ + +def softmax( + a: array, + /, + axis: int | Sequence[int] | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Perform the softmax along the given axis. + + This operation is a numerically stable version of: + + .. code-block:: + + exp(a) / sum(exp(a), axis, keepdims=True) + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or axes to compute + the softmax over. If unspecified this performs the softmax over + the full array. + + Returns: + array: The output of the softmax. + """ + +def sort( + a: array, + /, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Returns a sorted copy of the array. + + Args: + a (array): Input array. + axis (int or None, optional): Optional axis to sort over. + If ``None``, this sorts over the flattened array. + If unspecified, it defaults to -1 (sorting over the last axis). + + Returns: + array: The sorted array. + """ + +def split( + a: array, + /, + indices_or_sections: int | Sequence[int], + axis: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Split an array along a given axis. + + Args: + a (array): Input array. + indices_or_sections (int or list(int)): If ``indices_or_sections`` + is an integer the array is split into that many sections of equal + size. An error is raised if this is not possible. If ``indices_or_sections`` + is a list, the list contains the indices of the start of each subarray + along the given axis. + axis (int, optional): Axis to split along, defaults to `0`. + + Returns: + list(array): A list of split arrays. + """ + +def sqrt(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise square root. + + Args: + a (array): Input array. + + Returns: + array: The square root of ``a``. + """ + +def square(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise square. + + Args: + a (array): Input array. + + Returns: + array: The square of ``a``. + """ + +def squeeze( + a: array, + /, + axis: int | Sequence[int] | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Remove length one axes from an array. + + Args: + a (array): Input array. + axis (int or tuple(int), optional): Axes to remove. Defaults + to ``None`` in which case all size one axes are removed. + + Returns: + array: The output array with size one axes removed. + """ + +def stack( + arrays: list[array], + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Stacks the arrays along a new axis. + + Args: + arrays (list(array)): A list of arrays to stack. + axis (int, optional): The axis in the result array along which the + input arrays are stacked. Defaults to ``0``. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: The resulting stacked array. + """ + +def std( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + ddof: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Compute the standard deviation(s) over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + ddof (int, optional): The divisor to compute the variance + is ``N - ddof``, defaults to 0. + + Returns: + array: The output array of standard deviations. + """ + +def stop_gradient(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Stop gradients from being computed. + + The operation is the identity but it prevents gradients from flowing + through the array. + + Args: + a (array): Input array. + + Returns: + array: + The unchanged input ``a`` but without gradient flowing + through it. + """ + +def stream(s: Stream | Device) -> StreamContext: + """ + Create a context manager to set the default device and stream. + + Args: + s: The :obj:`Stream` or :obj:`Device` to set as the default. + + Returns: + A context manager that sets the default device and stream. + + Example: + + .. code-block::python + + import mlx.core as mx + + # Create a context manager for the default device and stream. + with mx.stream(mx.cpu): + # Operations here will use mx.cpu by default. + pass + """ + +def subtract( + a: scalar | array, + b: scalar | array, + stream: Stream | Device | None = ..., +) -> array: + """ + Element-wise subtraction. + + Subtract one array from another with numpy-style broadcasting semantics. Either or both + input arrays can also be scalars. + + Args: + a (array): Input array or scalar. + b (array): Input array or scalar. + + Returns: + array: The difference ``a - b``. + """ + +def sum( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Sum reduce the array over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + + Returns: + array: The output array with the corresponding axes reduced. + """ + +def swapaxes( + a: array, /, axis1: int, axis2: int, *, stream: Stream | Device | None = ... +) -> array: + """ + Swap two axes of an array. + + Args: + a (array): Input array. + axis1 (int): Specifies the first axis. + axis2 (int): Specifies the second axis. + + Returns: + array: The array with swapped axes. + """ + +def synchronize(stream: Stream | None = ...) -> None: + """ + Synchronize with the given stream. + + Args: + stream (Stream, optional): The stream to synchronize with. If ``None`` + then the default stream of the default device is used. + Default: ``None``. + """ + +def take( + a: array, + /, + indices: int | array, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Take elements along an axis. + + The elements are taken from ``indices`` along the specified axis. + If the axis is not specified the array is treated as a flattened + 1-D array prior to performing the take. + + As an example, if the ``axis=1`` this is equivalent to ``a[:, indices, ...]``. + + Args: + a (array): Input array. + indices (int or array): Integer index or input array with integral type. + axis (int, optional): Axis along which to perform the take. If unspecified + the array is treated as a flattened 1-D vector. + + Returns: + array: The indexed values of ``a``. + """ + +def take_along_axis( + a: array, + /, + indices: array, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Take values along an axis at the specified indices. + + Args: + a (array): Input array. + indices (array): Indices array. These should be broadcastable with + the input array excluding the `axis` dimension. + axis (int or None): Axis in the input to take the values from. If + ``axis == None`` the array is flattened to 1D prior to the indexing + operation. + + Returns: + array: The output array. + """ + +def tan(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise tangent. + + Args: + a (array): Input array. + + Returns: + array: The tangent of ``a``. + """ + +def tanh(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + Element-wise hyperbolic tangent. + + Args: + a (array): Input array. + + Returns: + array: The hyperbolic tangent of ``a``. + """ + +def tensordot( + a: array, + b: array, + /, + axes: int | list[Sequence[int]] = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Compute the tensor dot product along the specified axes. + + Args: + a (array): Input array + b (array): Input array + axes (int or list(list(int)), optional): The number of dimensions to + sum over. If an integer is provided, then sum over the last + ``axes`` dimensions of ``a`` and the first ``axes`` dimensions of + ``b``. If a list of lists is provided, then sum over the + corresponding dimensions of ``a`` and ``b``. Default: 2. + + Returns: + array: The tensor dot product. + """ + +def tile( + a: array, + reps: int | Sequence[int], + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Construct an array by repeating ``a`` the number of times given by ``reps``. + + Args: + a (array): Input array + reps (int or list(int)): The number of times to repeat ``a`` along each axis. + + Returns: + array: The tiled array. + """ + +def topk( + a: array, + /, + k: int, + axis: int | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Returns the ``k`` largest elements from the input along a given axis. + + The elements will not necessarily be in sorted order. + + Args: + a (array): Input array. + k (int): ``k`` top elements to be returned + axis (int or None, optional): Optional axis to select over. + If ``None``, this selects the top ``k`` elements over the + flattened array. If unspecified, it defaults to ``-1``. + + Returns: + array: The top ``k`` elements from the input. + """ + +def trace( + a: array, + /, + offset: int = ..., + axis1: int = ..., + axis2: int = ..., + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Return the sum along a specified diagonal in the given array. + + Args: + a (array): Input array + offset (int, optional): Offset of the diagonal from the main diagonal. + Can be positive or negative. Default: ``0``. + axis1 (int, optional): The first axis of the 2-D sub-arrays from which + the diagonals should be taken. Default: ``0``. + axis2 (int, optional): The second axis of the 2-D sub-arrays from which + the diagonals should be taken. Default: ``1``. + dtype (Dtype, optional): Data type of the output array. If + unspecified the output type is inferred from the input array. + + Returns: + array: Sum of specified diagonal. + """ + +def transpose( + a: array, + /, + axes: Sequence[int] | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Transpose the dimensions of the array. + + Args: + a (array): Input array. + axes (list(int), optional): Specifies the source axis for each axis + in the new array. The default is to reverse the axes. + + Returns: + array: The transposed array. + """ + +def tri( + n: int, + m: int, + k: int, + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + An array with ones at and below the given diagonal and zeros elsewhere. + + Args: + n (int): The number of rows in the output. + m (int, optional): The number of cols in the output. Defaults to ``None``. + k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. + dtype (Dtype, optional): Data type of the output array. Defaults to ``float32``. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: Array with its lower triangle filled with ones and zeros elsewhere + """ + +def tril(x: array, k: int, *, stream: Stream | Device | None = ...) -> array: + """ + Zeros the array above the given diagonal. + + Args: + x (array): input array. + k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: Array zeroed above the given diagonal + """ + +def triu(x: array, k: int, *, stream: Stream | Device | None = ...) -> array: + """ + Zeros the array below the given diagonal. + + Args: + x (array): input array. + k (int, optional): The diagonal of the 2-D array. Defaults to ``0``. + stream (Stream, optional): Stream or device. Defaults to ``None``. + + Returns: + array: Array zeroed below the given diagonal + """ + +uint16: Dtype = ... +uint32: Dtype = ... +uint64: Dtype = ... +uint8: Dtype = ... + +def unflatten( + a: array, + /, + axis: int, + shape: Sequence[int], + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Unflatten an axis of an array to a shape. + + Args: + a (array): Input array. + axis (int): The axis to unflatten. + shape (tuple(int)): The shape to unflatten to. At most one + entry can be ``-1`` in which case the corresponding size will be + inferred. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The unflattened array. + + Example: + >>> a = mx.array([1, 2, 3, 4]) + >>> mx.unflatten(a, 0, (2, -1)) + array([[1, 2], [3, 4]], dtype=int32) + """ + +unsignedinteger: DtypeCategory = ... + +def value_and_grad( + fun: Callable, + argnums: int | Sequence[int] | None = ..., + argnames: str | Sequence[str] = ..., +) -> Callable: + """ + Returns a function which computes the value and gradient of ``fun``. + + The function passed to :func:`value_and_grad` should return either + a scalar loss or a tuple in which the first element is a scalar + loss and the remaining elements can be anything. + + .. code-block:: python + + import mlx.core as mx + + def mse(params, inputs, targets): + outputs = forward(params, inputs) + lvalue = (outputs - targets).square().mean() + return lvalue + + # Returns lvalue, dlvalue/dparams + lvalue, grads = mx.value_and_grad(mse)(params, inputs, targets) + + def lasso(params, inputs, targets, a=1.0, b=1.0): + outputs = forward(params, inputs) + mse = (outputs - targets).square().mean() + l1 = mx.abs(outputs - targets).mean() + + loss = a*mse + b*l1 + + return loss, mse, l1 + + (loss, mse, l1), grads = mx.value_and_grad(lasso)(params, inputs, targets) + + Args: + fun (Callable): A function which takes a variable number of + :class:`array` or trees of :class:`array` and returns + a scalar output :class:`array` or a tuple the first element + of which should be a scalar :class:`array`. + argnums (int or list(int), optional): Specify the index (or indices) + of the positional arguments of ``fun`` to compute the gradient + with respect to. If neither ``argnums`` nor ``argnames`` are + provided ``argnums`` defaults to ``0`` indicating ``fun``'s first + argument. + argnames (str or list(str), optional): Specify keyword arguments of + ``fun`` to compute gradients with respect to. It defaults to [] so + no gradients for keyword arguments by default. + + Returns: + Callable: A function which returns a tuple where the first element + is the output of `fun` and the second element is the gradients w.r.t. + the loss. + """ + +def var( + a: array, + /, + axis: int | Sequence[int] | None = ..., + keepdims: bool = ..., + ddof: int = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Compute the variance(s) over the given axes. + + Args: + a (array): Input array. + axis (int or list(int), optional): Optional axis or + axes to reduce over. If unspecified this defaults + to reducing over the entire array. + keepdims (bool, optional): Keep reduced axes as + singleton dimensions, defaults to `False`. + ddof (int, optional): The divisor to compute the variance + is ``N - ddof``, defaults to 0. + + Returns: + array: The output array of variances. + """ + +def view( + a: scalar | array, dtype: Dtype, stream: Stream | Device | None = ... +) -> array: + """ + View the array as a different type. + + The output shape changes along the last axis if the input array's + type and the input ``dtype`` do not have the same size. + + Note: the view op does not imply that the input and output arrays share + their underlying data. The view only gaurantees that the binary + representation of each element (or group of elements) is the same. + + Args: + a (array): Input array or scalar. + dtype (Dtype): The data type to change to. + + Returns: + array: The array with the new type. + """ + +def vjp( + fun: Callable, primals: list[array], cotangents: list[array] +) -> tuple[list[array], list[array]]: + """ + Compute the vector-Jacobian product. + + Computes the product of the ``cotangents`` with the Jacobian of a + function ``fun`` evaluated at ``primals``. + + Args: + fun (Callable): A function which takes a variable number of :class:`array` + and returns a single :class:`array` or list of :class:`array`. + primals (list(array)): A list of :class:`array` at which to + evaluate the Jacobian. + cotangents (list(array)): A list of :class:`array` which are the + "vector" in the vector-Jacobian product. The ``cotangents`` should be the + same in number, shape, and type as the outputs of ``fun``. + + Returns: + list(array): A list of the vector-Jacobian products which + is the same in number, shape, and type of the outputs of ``fun``. + """ + +def vmap(fun: Callable, in_axes: object = ..., out_axes: object = ...) -> Callable: + """ + Returns a vectorized version of ``fun``. + + Args: + fun (Callable): A function which takes a variable number of + :class:`array` or a tree of :class:`array` and returns + a variable number of :class:`array` or a tree of :class:`array`. + in_axes (int, optional): An integer or a valid prefix tree of the + inputs to ``fun`` where each node specifies the vmapped axis. If + the value is ``None`` then the corresponding input(s) are not vmapped. + Defaults to ``0``. + out_axes (int, optional): An integer or a valid prefix tree of the + outputs of ``fun`` where each node specifies the vmapped axis. If + the value is ``None`` then the corresponding outputs(s) are not vmapped. + Defaults to ``0``. + + Returns: + Callable: The vectorized function. + """ + +def where( + condition: scalar | array, + x: scalar | array, + y: scalar | array, + /, + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Select from ``x`` or ``y`` according to ``condition``. + + The condition and input arrays must be the same shape or + broadcastable with each another. + + Args: + condition (array): The condition array. + x (array): The input selected from where condition is ``True``. + y (array): The input selected from where condition is ``False``. + + Returns: + array: The output containing elements selected from + ``x`` and ``y``. + """ + +def zeros( + shape: int | Sequence[int], + dtype: Dtype | None = ..., + *, + stream: Stream | Device | None = ..., +) -> array: + """ + Construct an array of zeros. + + Args: + shape (int or list(int)): The shape of the output array. + dtype (Dtype, optional): Data type of the output array. If + unspecified the output type defaults to ``float32``. + + Returns: + array: The array of zeros with the specified shape. + """ + +def zeros_like(a: array, /, *, stream: Stream | Device | None = ...) -> array: + """ + An array of zeros like the input. + + Args: + a (array): The input to take the shape and type from. + + Returns: + array: The output array filled with zeros. + """ + +scalar: TypeAlias = int | float | bool +list_or_scalar: TypeAlias = scalar | list["list_or_scalar"] +bool_: Dtype = ... diff --git a/typings/mlx/core/cuda/__init__.pyi b/typings/mlx/core/cuda/__init__.pyi new file mode 100644 index 00000000..cb7e23ba --- /dev/null +++ b/typings/mlx/core/cuda/__init__.pyi @@ -0,0 +1,2 @@ +def is_available() -> bool: + """Check if the CUDA back-end is available.""" diff --git a/typings/mlx/core/distributed/__init__.pyi b/typings/mlx/core/distributed/__init__.pyi new file mode 100644 index 00000000..15a952c4 --- /dev/null +++ b/typings/mlx/core/distributed/__init__.pyi @@ -0,0 +1,216 @@ +from typing import Sequence + +from mlx.core import Device, Dtype, Stream, array + +class Group: + """ + An :class:`mlx.core.distributed.Group` represents a group of independent mlx + processes that can communicate. + """ + def rank(self) -> int: + """Get the rank of this process""" + + def size(self) -> int: + """Get the size of the group""" + + def split(self, color: int, key: int = ...) -> Group: + """ + Split the group to subgroups based on the provided color. + + Processes that use the same color go to the same group. The ``key`` + argument defines the rank in the new group. The smaller the key the + smaller the rank. If the key is negative then the rank in the + current group is used. + + Args: + color (int): A value to group processes into subgroups. + key (int, optional): A key to optionally change the rank ordering + of the processes. + """ + +def all_gather( + x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... +) -> array: + """ + Gather arrays from all processes. + + Gather the ``x`` arrays from all processes in the group and concatenate + them along the first axis. The arrays should all have the same shape. + + Args: + x (array): Input array. + group (Group): The group of processes that will participate in the + gather. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The concatenation of all ``x`` arrays. + """ + +def all_max( + x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... +) -> array: + """ + All reduce max. + + Find the maximum of the ``x`` arrays from all processes in the group. + + Args: + x (array): Input array. + group (Group): The group of processes that will participate in the + reduction. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The maximum of all ``x`` arrays. + """ + +def all_min( + x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... +) -> array: + """ + All reduce min. + + Find the minimum of the ``x`` arrays from all processes in the group. + + Args: + x (array): Input array. + group (Group): The group of processes that will participate in the + reduction. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The minimum of all ``x`` arrays. + """ + +def all_sum( + x: array, *, group: Group | None = ..., stream: Stream | Device | None = ... +) -> array: + """ + All reduce sum. + + Sum the ``x`` arrays from all processes in the group. + + Args: + x (array): Input array. + group (Group): The group of processes that will participate in the + reduction. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The sum of all ``x`` arrays. + """ + +def init(strict: bool = ..., backend: str = ...) -> Group: + """ + Initialize the communication backend and create the global communication group. + + Example: + + .. code:: python + + import mlx.core as mx + + group = mx.distributed.init(backend="ring") + + Args: + strict (bool, optional): If set to False it returns a singleton group + in case ``mx.distributed.is_available()`` returns False otherwise + it throws a runtime error. Default: ``False`` + backend (str, optional): Which distributed backend to initialize. + Possible values ``mpi``, ``ring``, ``nccl``, ``any``. If set to ``any`` all + available backends are tried and the first one that succeeds + becomes the global group which will be returned in subsequent + calls. Default: ``any`` + + Returns: + Group: The group representing all the launched processes. + """ + +def is_available() -> bool: + """Check if a communication backend is available.""" + +def recv( + shape: Sequence[int], + dtype: Dtype, + src: int, + *, + group: Group | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Recv an array with shape ``shape`` and dtype ``dtype`` from process + with rank ``src``. + + Args: + shape (tuple[int]): The shape of the array we are receiving. + dtype (Dtype): The data type of the array we are receiving. + src (int): Rank of the source process in the group. + group (Group): The group of processes that will participate in the + recv. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The array that was received from ``src``. + """ + +def recv_like( + x: array, + src: int, + *, + group: Group | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Recv an array with shape and type like ``x`` from process with rank + ``src``. + + It is equivalent to calling ``mx.distributed.recv(x.shape, x.dtype, src)``. + + Args: + x (array): An array defining the shape and dtype of the array we are + receiving. + src (int): Rank of the source process in the group. + group (Group): The group of processes that will participate in the + recv. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The array that was received from ``src``. + """ + +def send( + x: array, + dst: int, + *, + group: Group | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Send an array from the current process to the process that has rank + ``dst`` in the group. + + Args: + x (array): Input array. + dst (int): Rank of the destination process in the group. + group (Group): The group of processes that will participate in the + sned. If set to ``None`` the global group is used. Default: + ``None``. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: An array identical to ``x`` which when evaluated the send is performed. + """ diff --git a/typings/mlx/core/metal/__init__.pyi b/typings/mlx/core/metal/__init__.pyi new file mode 100644 index 00000000..983f0067 --- /dev/null +++ b/typings/mlx/core/metal/__init__.pyi @@ -0,0 +1,38 @@ +def clear_cache() -> None: ... +def device_info() -> dict[str, str | int]: + """ + Get information about the GPU device and system settings. + + Currently returns: + + * ``architecture`` + * ``max_buffer_size`` + * ``max_recommended_working_set_size`` + * ``memory_size`` + * ``resource_limit`` + + Returns: + dict: A dictionary with string keys and string or integer values. + """ + +def get_active_memory() -> int: ... +def get_cache_memory() -> int: ... +def get_peak_memory() -> int: ... +def is_available() -> bool: + """Check if the Metal back-end is available.""" + +def reset_peak_memory() -> None: ... +def set_cache_limit(limit: int) -> int: ... +def set_memory_limit(limit: int) -> int: ... +def set_wired_limit(limit: int) -> int: ... +def start_capture(path: str) -> None: + """ + Start a Metal capture. + + Args: + path (str): The path to save the capture which should have + the extension ``.gputrace``. + """ + +def stop_capture() -> None: + """Stop a Metal capture.""" diff --git a/typings/mlx/core/random/__init__.pyi b/typings/mlx/core/random/__init__.pyi new file mode 100644 index 00000000..4116e0ec --- /dev/null +++ b/typings/mlx/core/random/__init__.pyi @@ -0,0 +1,301 @@ +from typing import Sequence + +from mlx.core import Device, Dtype, Stream, array, scalar +from mlx.core.distributed import state as state + +def bernoulli( + p: scalar | array = ..., + shape: Sequence[int] | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate Bernoulli random values. + + The values are sampled from the bernoulli distribution with parameter + ``p``. The parameter ``p`` can be a :obj:`float` or :obj:`array` and + must be broadcastable to ``shape``. + + Args: + p (float or array, optional): Parameter of the Bernoulli + distribution. Default: ``0.5``. + shape (list(int), optional): Shape of the output. + Default: ``p.shape``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The array of random integers. + """ + +def categorical( + logits: array, + axis: int = ..., + shape: Sequence[int] | None = ..., + num_samples: int | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Sample from a categorical distribution. + + The values are sampled from the categorical distribution specified by + the unnormalized values in ``logits``. Note, at most one of ``shape`` + or ``num_samples`` can be specified. If both are ``None``, the output + has the same shape as ``logits`` with the ``axis`` dimension removed. + + Args: + logits (array): The *unnormalized* categorical distribution(s). + axis (int, optional): The axis which specifies the distribution. + Default: ``-1``. + shape (list(int), optional): The shape of the output. This must + be broadcast compatible with ``logits.shape`` with the ``axis`` + dimension removed. Default: ``None`` + num_samples (int, optional): The number of samples to draw from each + of the categorical distributions in ``logits``. The output will have + ``num_samples`` in the last dimension. Default: ``None``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The ``shape``-sized output array with type ``uint32``. + """ + +def gumbel( + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + key: Stream | Device | None = ..., + stream: array | None = ..., +) -> array: + """ + Sample from the standard Gumbel distribution. + + The values are sampled from a standard Gumbel distribution + which CDF ``exp(-exp(-x))``. + + Args: + shape (list(int)): The shape of the output. + dtype (Dtype, optional): The data type of the output. + Default: ``float32``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: + The :class:`array` with shape ``shape`` and distributed according + to the Gumbel distribution. + """ + +def key(seed: int) -> array: + """ + Get a PRNG key from a seed. + + Args: + seed (int): Seed for the PRNG. + + Returns: + array: The PRNG key array. + """ + +def laplace( + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + loc: float = ..., + scale: float = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Sample numbers from a Laplace distribution. + + Args: + shape (list(int), optional): Shape of the output. Default: ``()``. + dtype (Dtype, optional): Type of the output. Default: ``float32``. + loc (float, optional): Mean of the distribution. Default: ``0.0``. + scale (float, optional): The scale "b" of the Laplace distribution. + Default:``1.0``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The output array of random values. + """ + +def multivariate_normal( + mean: array, + cov: array, + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate jointly-normal random samples given a mean and covariance. + + The matrix ``cov`` must be positive semi-definite. The behavior is + undefined if it is not. The only supported ``dtype`` is ``float32``. + + Args: + mean (array): array of shape ``(..., n)``, the mean of the + distribution. + cov (array): array of shape ``(..., n, n)``, the covariance + matrix of the distribution. The batch shape ``...`` must be + broadcast-compatible with that of ``mean``. + shape (list(int), optional): The output shape must be + broadcast-compatible with ``mean.shape[:-1]`` and ``cov.shape[:-2]``. + If empty, the result shape is determined by broadcasting the batch + shapes of ``mean`` and ``cov``. Default: ``[]``. + dtype (Dtype, optional): The output type. Default: ``float32``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The output array of random values. + """ + +def normal( + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + loc: scalar | array | None = ..., + scale: scalar | array | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + r""" + Generate normally distributed random numbers. + + If ``loc`` and ``scale`` are not provided the "standard" normal + distribution is used. That means $x \sim \mathcal{N}(0, 1)$ for + real numbers and $\text{Re}(x),\text{Im}(x) \sim \mathcal{N}(0, + \frac{1}{2})$ for complex numbers. + + Args: + shape (list(int), optional): Shape of the output. Default: ``()``. + dtype (Dtype, optional): Type of the output. Default: ``float32``. + loc (scalar or array, optional): Mean of the distribution. + Default: ``None``. + scale (scalar or array, optional): Standard deviation of the + distribution. Default: ``None``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The output array of random values. + """ + +def permutation( + x: int | array, + axis: int = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate a random permutation or permute the entries of an array. + + Args: + x (int or array, optional): If an integer is provided a random + permtuation of ``mx.arange(x)`` is returned. Otherwise the entries + of ``x`` along the given axis are randomly permuted. + axis (int, optional): The axis to permute along. Default: ``0``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: + The generated random permutation or randomly permuted input array. + """ + +def randint( + low: scalar | array, + high: scalar | array, + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate random integers from the given interval. + + The values are sampled with equal probability from the integers in + half-open interval ``[low, high)``. The lower and upper bound can be + scalars or arrays and must be broadcastable to ``shape``. + + Args: + low (scalar or array): Lower bound of the interval. + high (scalar or array): Upper bound of the interval. + shape (list(int), optional): Shape of the output. Default: ``()``. + dtype (Dtype, optional): Type of the output. Default: ``int32``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The array of random integers. + """ + +def seed(seed: int) -> None: + """ + Seed the global PRNG. + + Args: + seed (int): Seed for the global PRNG. + """ + +def split(key: array, num: int = ..., stream: Stream | Device | None = ...) -> array: + """ + Split a PRNG key into sub keys. + + Args: + key (array): Input key to split. + num (int, optional): Number of sub keys. Default: ``2``. + + Returns: + array: The array of sub keys with ``num`` as its first dimension. + """ + +def truncated_normal( + lower: scalar | array, + upper: scalar | array, + shape: Sequence[int] | None = ..., + dtype: Dtype | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate values from a truncated normal distribution. + + The values are sampled from the truncated normal distribution + on the domain ``(lower, upper)``. The bounds ``lower`` and ``upper`` + can be scalars or arrays and must be broadcastable to ``shape``. + + Args: + lower (scalar or array): Lower bound of the domain. + upper (scalar or array): Upper bound of the domain. + shape (list(int), optional): The shape of the output. + Default:``()``. + dtype (Dtype, optional): The data type of the output. + Default: ``float32``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The output array of random values. + """ + +def uniform( + low: scalar | array = ..., + high: scalar | array = ..., + shape: Sequence[int] = ..., + dtype: Dtype | None = ..., + key: array | None = ..., + stream: Stream | Device | None = ..., +) -> array: + """ + Generate uniformly distributed random numbers. + + The values are sampled uniformly in the half-open interval ``[low, high)``. + The lower and upper bound can be scalars or arrays and must be + broadcastable to ``shape``. + + Args: + low (scalar or array, optional): Lower bound of the distribution. + Default: ``0``. + high (scalar or array, optional): Upper bound of the distribution. + Default: ``1``. + shape (list(int), optional): Shape of the output. Default:``()``. + dtype (Dtype, optional): Type of the output. Default: ``float32``. + key (array, optional): A PRNG key. Default: ``None``. + + Returns: + array: The output array random values. + """ diff --git a/uv.lock b/uv.lock index 426cfd70..deabdc7b 100644 --- a/uv.lock +++ b/uv.lock @@ -14,7 +14,6 @@ supported-markers = [ members = [ "exo", "exo-pyo3-bindings", - "exo-scripts", ] [[package]] @@ -438,21 +437,6 @@ dev = [ { name = "pytest-asyncio", specifier = ">=1.0.0" }, ] -[[package]] -name = "exo-scripts" -version = "0.1.0" -source = { editable = "scripts" } -dependencies = [ - { name = "exo", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, -] - -[package.metadata] -requires-dist = [ - { name = "exo", editable = "." }, - { name = "huggingface-hub", specifier = ">=0.33.4" }, -] - [[package]] name = "fastapi" version = "0.121.0" @@ -561,12 +545,16 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/43/3cecdc0349359e1a527cbf2e3e28e5f8f06d3343aaf82ca13437a9aa290f/greenlet-3.2.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23768528f2911bcd7e475210822ffb5254ed10d71f4028387e5a99b4c6699671", size = 610497, upload-time = "2025-08-07T13:18:31.636Z" }, { url = "https://files.pythonhosted.org/packages/b8/19/06b6cf5d604e2c382a6f31cafafd6f33d5dea706f4db7bdab184bad2b21d/greenlet-3.2.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:00fadb3fedccc447f517ee0d3fd8fe49eae949e1cd0f6a611818f4f6fb7dc83b", size = 1121662, upload-time = "2025-08-07T13:42:41.117Z" }, { url = "https://files.pythonhosted.org/packages/a2/15/0d5e4e1a66fab130d98168fe984c509249c833c1a3c16806b90f253ce7b9/greenlet-3.2.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:d25c5091190f2dc0eaa3f950252122edbbadbb682aa7b1ef2f8af0f8c0afefae", size = 1149210, upload-time = "2025-08-07T13:18:24.072Z" }, + { url = "https://files.pythonhosted.org/packages/1c/53/f9c440463b3057485b8594d7a638bed53ba531165ef0ca0e6c364b5cc807/greenlet-3.2.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6e343822feb58ac4d0a1211bd9399de2b3a04963ddeec21530fc426cc121f19b", size = 1564759, upload-time = "2025-11-04T12:42:19.395Z" }, + { url = "https://files.pythonhosted.org/packages/47/e4/3bb4240abdd0a8d23f4f88adec746a3099f0d86bfedb623f063b2e3b4df0/greenlet-3.2.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca7f6f1f2649b89ce02f6f229d7c19f680a6238af656f61e0115b24857917929", size = 1634288, upload-time = "2025-11-04T12:42:21.174Z" }, { url = "https://files.pythonhosted.org/packages/22/5c/85273fd7cc388285632b0498dbbab97596e04b154933dfe0f3e68156c68c/greenlet-3.2.4-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:49a30d5fda2507ae77be16479bdb62a660fa51b1eb4928b524975b3bde77b3c0", size = 273586, upload-time = "2025-08-07T13:16:08.004Z" }, { url = "https://files.pythonhosted.org/packages/d1/75/10aeeaa3da9332c2e761e4c50d4c3556c21113ee3f0afa2cf5769946f7a3/greenlet-3.2.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:299fd615cd8fc86267b47597123e3f43ad79c9d8a22bebdce535e53550763e2f", size = 686346, upload-time = "2025-08-07T13:42:59.944Z" }, { url = "https://files.pythonhosted.org/packages/c0/aa/687d6b12ffb505a4447567d1f3abea23bd20e73a5bed63871178e0831b7a/greenlet-3.2.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:c17b6b34111ea72fc5a4e4beec9711d2226285f0386ea83477cbb97c30a3f3a5", size = 699218, upload-time = "2025-08-07T13:45:30.969Z" }, { url = "https://files.pythonhosted.org/packages/dc/8b/29aae55436521f1d6f8ff4e12fb676f3400de7fcf27fccd1d4d17fd8fecd/greenlet-3.2.4-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b4a1870c51720687af7fa3e7cda6d08d801dae660f75a76f3845b642b4da6ee1", size = 694659, upload-time = "2025-08-07T13:53:17.759Z" }, { url = "https://files.pythonhosted.org/packages/92/2e/ea25914b1ebfde93b6fc4ff46d6864564fba59024e928bdc7de475affc25/greenlet-3.2.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:061dc4cf2c34852b052a8620d40f36324554bc192be474b9e9770e8c042fd735", size = 695355, upload-time = "2025-08-07T13:18:34.517Z" }, { url = "https://files.pythonhosted.org/packages/72/60/fc56c62046ec17f6b0d3060564562c64c862948c9d4bc8aa807cf5bd74f4/greenlet-3.2.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44358b9bf66c8576a9f57a590d5f5d6e72fa4228b763d0e43fee6d3b06d3a337", size = 657512, upload-time = "2025-08-07T13:18:33.969Z" }, + { url = "https://files.pythonhosted.org/packages/23/6e/74407aed965a4ab6ddd93a7ded3180b730d281c77b765788419484cdfeef/greenlet-3.2.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2917bdf657f5859fbf3386b12d68ede4cf1f04c90c3a6bc1f013dd68a22e2269", size = 1612508, upload-time = "2025-11-04T12:42:23.427Z" }, + { url = "https://files.pythonhosted.org/packages/0d/da/343cd760ab2f92bac1845ca07ee3faea9fe52bee65f7bcb19f16ad7de08b/greenlet-3.2.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:015d48959d4add5d6c9f6c5210ee3803a830dce46356e3bc326d6776bde54681", size = 1680760, upload-time = "2025-11-04T12:42:25.341Z" }, ] [[package]]