"""mlx.core.fast: fast operations""" from collections.abc import Sequence from typing import Annotated from numpy.typing import NDArray import mlx.core from mlx.core import array, Dtype, Device, Stream, scalar from typing import Sequence, Optional, Union def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array: """ Root Mean Square normalization (RMS norm). The normalization is with respect to the last axis of the input ``x``. Args: x (array): Input array. weight (array, optional): A multiplicative weight to scale the result by. The ``weight`` should be one-dimensional with the same size as the last axis of ``x``. If set to ``None`` then no scaling happens. eps (float): A small additive constant for numerical stability. Returns: array: The output array. """ def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array: """ Layer normalization. The normalization is with respect to the last axis of the input ``x``. Args: x (array): Input array. weight (array, optional): A multiplicative weight to scale the result by. The ``weight`` should be one-dimensional with the same size as the last axis of ``x``. If set to ``None`` then no scaling happens. bias (array, optional): An additive offset to be added to the result. The ``bias`` should be one-dimensional with the same size as the last axis of ``x``. If set to ``None`` then no translation happens. eps (float): A small additive constant for numerical stability. Returns: array: The output array. """ def rope(a: array, dims: int, *, traditional: bool, base: Optional[float], scale: float, offset: Union[int, array], freqs: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array: """ Apply rotary positional encoding to the input. The input is expected to be at least 3D with shape ``(B, *, T, D)`` where: * ``B`` is the batch size. * ``T`` is the sequence length. * ``D`` is the feature dimension. Args: a (array): The input array. dims (int): The feature dimensions to be rotated. If the input feature is larger than dims then the rest is left unchanged. traditional (bool): If set to ``True`` choose the traditional implementation which rotates consecutive dimensions. base (float, optional): The base used to compute angular frequency for each dimension in the positional encodings. Exactly one of ``base`` and ``freqs`` must be ``None``. scale (float): The scale used to scale the positions. offset (int or array): The position offset to start at. If an :obj:`array` is given it can be a scalar or vector of ``B`` offsets for each example in the batch. freqs (array, optional): Optional frequencies to use with RoPE. If set, the ``base`` parameter must be ``None``. Default: ``None``. Returns: array: The output array. """ def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: Union[None, str, array] = None, sinks: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array: """ A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``. Supports: * `Multi-Head Attention `_ * `Grouped Query Attention `_ * `Multi-Query Attention `_ .. note:: * The softmax operation is performed in ``float32`` regardless of the input precision. * For Grouped Query Attention and Multi-Query Attention, the ``k`` and ``v`` inputs should not be pre-tiled to match ``q``. In the following the dimensions are given by: * ``B``: The batch size. * ``N_q``: The number of query heads. * ``N_kv``: The number of key and value heads. * ``T_q``: The number of queries per example. * ``T_kv``: The number of keys and values per example. * ``D``: The per-head dimension. Args: q (array): Queries with shape ``[B, N_q, T_q, D]``. k (array): Keys with shape ``[B, N_kv, T_kv, D]``. v (array): Values with shape ``[B, N_kv, T_kv, D]``. scale (float): Scale for queries (typically ``1.0 / sqrt(q.shape(-1)``). mask (str or array, optional): The mask to apply to the query-key scores. The mask can be an array or a string indicating the mask type. The only supported string type is ``"causal"``. If the mask is an array it can be a boolean or additive mask. The mask can have at most 4 dimensions and must be broadcast-compatible with the shape ``[B, N, T_q, T_kv]``. If an additive mask is given its type must promote to the promoted type of ``q``, ``k``, and ``v``. The ``"causal"`` mask uses lower-right alignment where the last query aligns with the last key. sinks (array, optional): An optional array of attention sinks. Default: ``None``. Returns: array: The output array. Example: .. code-block:: python B = 2 N_q = N_kv = 32 T_q = T_kv = 1000 D = 128 q = mx.random.normal(shape=(B, N_q, T_q, D)) k = mx.random.normal(shape=(B, N_kv, T_kv, D)) v = mx.random.normal(shape=(B, N_kv, T_kv, D)) scale = D ** -0.5 out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask="causal") """ def metal_kernel(name: str, input_names: Sequence[str], output_names: Sequence[str], source: str, header: str = '', ensure_row_contiguous: bool = True, atomic_outputs: bool = False) -> object: r""" A jit-compiled custom Metal kernel defined from a source string. Full documentation: :ref:`custom_metal_kernels`. Args: name (str): Name for the kernel. input_names (List[str]): The parameter names of the inputs in the function signature. output_names (List[str]): The parameter names of the outputs in the function signature. source (str): Source code. This is the body of a function in Metal, the function signature will be automatically generated. header (str): Header source code to include before the main function. Useful for helper functions or includes that should live outside of the main function body. ensure_row_contiguous (bool): Whether to ensure the inputs are row contiguous before the kernel runs. Default: ``True``. atomic_outputs (bool): Whether to use atomic outputs in the function signature e.g. ``device atomic``. Default: ``False``. Returns: Callable ``metal_kernel``. Example: .. code-block:: python def exp_elementwise(a: mx.array): source = \'\'' uint elem = thread_position_in_grid.x; T tmp = inp[elem]; out[elem] = metal::exp(tmp); \'\'' kernel = mx.fast.metal_kernel( name="myexp", input_names=["inp"], output_names=["out"], source=source ) outputs = kernel( inputs=[a], template=[("T", mx.float32)], grid=(a.size, 1, 1), threadgroup=(256, 1, 1), output_shapes=[a.shape], output_dtypes=[a.dtype], verbose=True, ) return outputs[0] a = mx.random.normal(shape=(4, 16)).astype(mx.float16) b = exp_elementwise(a) assert mx.allclose(b, mx.exp(a)) """ def cuda_kernel(name: str, input_names: Sequence[str], output_names: Sequence[str], source: str, header: str = '', ensure_row_contiguous: bool = True, shared_memory: int = 0) -> object: r""" A jit-compiled custom CUDA kernel defined from a source string. This is the CUDA equivalent of :ref:`custom_metal_kernels`. Args: name (str): Name for the kernel. input_names (List[str]): The parameter names of the inputs in the function signature. output_names (List[str]): The parameter names of the outputs in the function signature. source (str): Source code. This is the body of a function in CUDA, the function signature will be automatically generated. header (str): Header source code to include before the main function. Useful for helper functions or includes that should live outside of the main function body. ensure_row_contiguous (bool): Whether to ensure the inputs are row contiguous before the kernel runs. Default: ``True``. shared_memory (int): The dynamic shared memory to request for the kernel. A value of 0 means no dynamic shared memory. Default: ``0``. Returns: Callable ``cuda_kernel``. Example: .. code-block:: python def exp_elementwise(a: mx.array): source = \'\'' auto elem = cooperative_groups::this_grid().thread_rank(); T tmp = inp[elem]; out[elem] = exp(tmp); \'\'' kernel = mx.fast.cuda_kernel( name="myexp", input_names=["inp"], output_names=["out"], source=source ) outputs = kernel( inputs=[a], template=[("T", mx.float32)], grid=(a.size, 1, 1), threadgroup=(256, 1, 1), output_shapes=[a.shape], output_dtypes=[a.dtype], verbose=True, ) return outputs[0] a = mx.random.normal(shape=(16, 16)).astype(mx.float16) b = exp_elementwise(a) assert mx.allclose(b, mx.exp(a)) """ def precompiled_cuda_kernel(*, name: str, compiled_source: bytes, inputs: Sequence[bool | int | float | mlx.core.array | Annotated[NDArray, dict(order='C', device='cpu', writable=False)] | complex | mlx.core.ArrayLike], output_shapes: Sequence[tuple[int, ...]], output_dtypes: Sequence[mlx.core.Dtype], scalars: Sequence[object], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], shared_memory: int = 0, init_value: float | None = None, ensure_row_contiguous: bool = False, stream: mlx.core.Stream | mlx.core.ThreadLocalStream | mlx.core.Device | None = None) -> list[mlx.core.array]: """ Run a precompiled CUDA kernel defined from PTX or cubin. This op is still experimental and various parts of the API may change. Args: name (str): Name for the kernel compiled_source (bytes): The precompiled kernel in raw bytes. inputs (List[array]): The inputs passed to the CUDA kernel. output_shapes (List[Sequence[int]]): The list of shapes for each output. output_dtypes (List[Dtype]): The list of data types for each output. scalars (List[Union[bool, int, float]]): A list of scalar arguments to pass to the kernel. grid (tuple[int, int, int]): 3-tuple specifying the grid to launch the kernel with. For compatibility with :func:`metal_kernel` the grid is in threads and not in threadblocks. threadgroup (tuple[int, int, int]): 3-tuple specifying the threadgroup size to use. shared_memory (int): The dynamic shared memory to request for the kernel. A value of 0 means no dynamic shared memory. Default: ``0``. init_value (float, optional): Optional value to use to initialize all of the output arrays. By default, output arrays are uninitialized. Default: ``None``. ensure_row_contiguous (bool): Whether to ensure the inputs are row contiguous before the kernel runs. Default: ``False``. stream (mx.stream, optional): Stream to run the kernel on. Default: ``None``. """