anima_mlx/models/primitives.py

"""Small MLX primitives shared by Anima model components."""

from __future__ import annotations

import os
from typing import Any


USE_FAST_NORMS = os.environ.get("ANIMA_MLX_FAST_NORMS") == "1"


def linear(x: Any, weight: Any, bias: Any | None = None) -> Any:
    """Apply a PyTorch-layout linear weight to the last dimension of ``x``."""

    import mlx.core as mx

    y = x @ mx.transpose(weight)
    if bias is not None:
        y = y + bias
    return y


def rms_norm(x: Any, weight: Any, eps: float = 1e-6) -> Any:
    import mlx.core as mx

    if USE_FAST_NORMS:
        output = mx.fast.rms_norm(x, weight=weight.astype(x.dtype), eps=eps)
        if x.dtype in (mx.float16, mx.bfloat16):
            return output.astype(x.dtype)
        return output
    x32 = x.astype(mx.float32)
    y = x32 * mx.rsqrt(mx.mean(mx.square(x32), axis=-1, keepdims=True) + eps)
    output = y * weight.astype(mx.float32)
    if x.dtype in (mx.float16, mx.bfloat16):
        return output.astype(x.dtype)
    return output


def layer_norm(x: Any, eps: float = 1e-6) -> Any:
    import mlx.core as mx

    if USE_FAST_NORMS:
        weight = mx.ones((x.shape[-1],), dtype=x.dtype)
        bias = mx.zeros((x.shape[-1],), dtype=x.dtype)
        output = mx.fast.layer_norm(x, weight=weight, bias=bias, eps=eps)
        if x.dtype in (mx.float16, mx.bfloat16):
            return output.astype(x.dtype)
        return output
    x32 = x.astype(mx.float32)
    mean = mx.mean(x32, axis=-1, keepdims=True)
    variance = mx.mean(mx.square(x32 - mean), axis=-1, keepdims=True)
    output = (x32 - mean) * mx.rsqrt(variance + eps)
    if x.dtype in (mx.float16, mx.bfloat16):
        return output.astype(x.dtype)
    return output


def gelu(x: Any) -> Any:
    import mlx.core as mx

    return 0.5 * x * (1.0 + mx.erf(x / mx.sqrt(mx.array(2.0, dtype=x.dtype))))


def silu(x: Any) -> Any:
    import mlx.core as mx

    if x.dtype in (mx.float16, mx.bfloat16):
        x32 = x.astype(mx.float32)
        return (x32 * mx.sigmoid(x32)).astype(x.dtype)
    return x * mx.sigmoid(x)


def rotate_half(x: Any) -> Any:
    import mlx.core as mx

    half = x.shape[-1] // 2
    x1 = x[..., :half]
    x2 = x[..., half:]
    return mx.concatenate([-x2, x1], axis=-1)


def apply_rotary_pos_emb(x: Any, cos: Any, sin: Any, unsqueeze_dim: int = 1) -> Any:
    import mlx.core as mx

    cos = mx.expand_dims(cos, axis=unsqueeze_dim)
    sin = mx.expand_dims(sin, axis=unsqueeze_dim)
    return (x * cos) + (rotate_half(x) * sin)


def scaled_dot_product_attention(q: Any, k: Any, v: Any) -> Any:
    from .attention import scaled_dot_product_attention as attention

    return attention(q, k, v)