# Apple MLX: потеря масштабной инвариантности производных arctan2 Локальный технический отчёт, 10 сентября 2026 года. Дефект воспроизведён на MLX 0.32.2 и настоящем C++ runtime, CPU. Исследовательский патч прошёл 1368 сравнений; неуспешные промежуточные варианты сохранены. Это один дефект вычисления производных arctan2, а не 772 самостоятельные находки. ## Главный контрпример Для любого C>0 функция f(t) = atan2(C·t, C) = atan(t) не зависит от C. При t=1 её первые три производные в точности равны 0.5, −0.5, 0.5. | dtype | C | MLX: f′(1) | MLX: f″(1) | MLX: f‴(1) | Центральная разность forward | |---|---:|---:|---:|---:|---:| | float32 | 1 | 0.5 | −0.5 | 0.5 | 0.5000209808 | | float32 | 2⁻⁸⁰ | **inf** | NaN | NaN | 0.5000209808 | | float32 | 2⁸⁰ | **0** | NaN | NaN | 0.5000209808 | | float64 | 1 | 0.5 | −0.5 | 0.5 | 0.5000000795 | | float64 | 2⁻⁶⁰⁰ | **inf** | NaN | NaN | 0.5000000795 | | float64 | 2⁶⁰⁰ | **0** | NaN | NaN | 0.5000000795 | Прямой результат MLX во всех строках корректен с точностью dtype: около π/4. Центральные разности вычислены по значениям **настоящего MLX forward**, с шагом 2⁻⁶ для float32 и 2⁻¹⁰ для float64. Они независимо подтверждают первый градиент; второй и третий порядки проверяются аналитическими производными atan(t). Минимальное воспроизведение: import mlx.core as mx mx.set_default_device(mx.cpu) t = mx.array(1., dtype=mx.float32) for power in (0, -80, 80): c = mx.array(2.**power, dtype=mx.float32) f = lambda z: mx.arctan2(c*z, c) print(power, f(t).item(), mx.grad(f)(t).item()) # одинаковый forward, но градиенты: 0.5, inf, 0 [Python-пробник](probe.py), [все результаты wheel](wheel-reproduction.json). Результат не объясняется особой точкой (0,0): аргументы ненулевые, а составная функция равна гладкой atan(t) независимо от C. ## Первый порядок и геометрические инварианты В точке (y,x)=(C,C) градиент atan2(y,x) равен [1/(2C),−1/(2C)]. JVP в радиальном направлении (C,C) должен быть 0, в направлении (C,−C) — 1. В float16 при C=1024 исходный VJP возвращает [0,−0] вместо [0.00048828125,−0.00048828125]. При C=2⁻¹⁴ возвращается [inf,−inf] вместо [8192,−8192]. Оба ожидаемых результата представлены в нормальном диапазоне float16. Радиальный JVP в обоих случаях равен NaN вместо 0. Аналогичные нарушения воспроизведены для bfloat16, float32 и float64. Для обычных эталонов применяется std::hypot(y,x), затем (x/r)/r и −(y/r)/r; квадрат большого или малого радиуса не формируется. Для отдельных случаев с очень разными аргументами и весами эталон вычисляется непосредственно как степень двойки. Пренебрегаемая поправка квадрата отношения меньше точности эталона и допуска. ## Причина и вариант исправления Публичные [ArcTan2::vjp](https://github.com/ml-explore/mlx/blob/81ba1c6a0e50a9268b931579c2d4f1158b9aab5a/mlx/primitives.cpp#L524) и [ArcTan2::jvp](https://github.com/ml-explore/mlx/blob/81ba1c6a0e50a9268b931579c2d4f1158b9aab5a/mlx/primitives.cpp#L552) используют y²+x² непосредственно. Этот знаменатель может стать inf или 0 при конечных аргументах и представимых производных. JVP дополнительно формирует ненормированные произведения координат и направлений. [Локальный патч](arctan2-scale-autodiff.patch) выбирает s=max(|y|,|x|) и D=(y/s)²+(x/s)². При конечной ненулевой паре 1≤D≤2. Для каждого вклада вычисляется g·z/(s²D), где g — входящий вес, z=x для производной по y и z=−y для производной по x. Порядок операций выбирается по масштабу и представимости g/s: | Условие | Вычисление вклада | |---|---| | s≥1 | ((g/s)·z/D)/s | | 0