Гайд · TNWS AI
Как сделать условную ветку под jit через jax.lax.cond
Динамическая scalar-ветка в JAX 0.11.2: lax.cond, одинаковые output contracts, две фактически проверенные ветви под jit.
Задача и применимость
Актуальный запрос: «как использовать jax lax cond внутри jit». Выбрать одну из двух ветвей по scalar predicate, который зависит от runtime-значения массива внутри jitted функции. Обычный Python if не подходит для traced boolean такого типа. jax.lax.cond(pred, true_fun, false_fun, operand) требует совместимых структур, форм и dtype результатов обеих ветвей.
Материал применим к JAX 0.11.2 и jaxlib 0.11.2. JAX предоставляет массивы, автоматическое дифференцирование и программные преобразования; он не является готовой библиотекой слоёв, датасетов или полного тренировочного цикла. Пример намеренно мал и выполняется локально на доступном backend без загрузки весов и внешних данных. Численный критерий доказывает поведение конкретного кода, но не качество рабочей нейросети.
Пошаговое решение
- Вычислите scalar predicate из входного массива.
- Опишите true_fun и false_fun как чистые функции одного operand.
- Вызовите lax.cond внутри jax.jit.
- Проверьте положительный и отрицательный inputs, чтобы фактически пройти обе ветви.
Готовый проверяемый шаблон
import jax
import jax.numpy as jnp
@jax.jit
def route(x):
pred = jnp.mean(x) > 0
return jax.lax.cond(pred, lambda z: z * 2., lambda z: -z, x)
positive = route(jnp.array([1., 2.]))
negative = route(jnp.array([-1., -2.]))
assert jnp.array_equal(positive, jnp.array([2., 4.]))
assert jnp.array_equal(negative, jnp.array([1., 2.]))
assert positive.shape == negative.shape == (2,)
print(positive, negative)
Вход и ожидаемый результат
Для среднего больше нуля true-ветка удваивает [1,2] до [2,4]. Для отрицательного среднего false-ветка меняет знак [-1,-2] и даёт [1,2]. Обе ветви возвращают float-массив формы (2,).
Как проверить результат
Добавьте тест predicate на границе, проверьте обе ветви без jit и под jit, затем сравните shapes/dtypes. Учтите, что при преобразовании обе функции трассируются; не помещайте туда небезопасные побочные эффекты.
Перенос на настоящую нейросеть
Для рабочей модели сначала зафиксируйте контракт параметров и batch: pytree structure, формы, dtype, порядок признаков и семантику выходов. Генерацию случайностей оформите как явный поток keys. Разделите чистые функции forward, loss и update, чтобы их можно было независимо тестировать, дифференцировать и компилировать. Данные train, validation и test не должны смешиваться. Любое изменение формы, precision, backend или версии JAX проверяйте на сохранённом обезличенном regression-наборе до развёртывания.
Критерии приёмки
- Все формы и dtype совпадают с контрактом.
- Каждый PRNG subkey используется в одной логической случайной операции.
- Loss и gradients конечны на положительном тесте.
- Негативный вход вызывает ожидаемую ошибку, а не случайный output.
- Обычный и преобразованный через jit/vmap путь численно согласованы, когда это применимо.
- Устройство, версии JAX/jaxlib и численный допуск записаны в отчёте.
Регрессионный протокол
Храните несколько малых входов: стандартный batch, один объект, нули, граничные значения и заведомо невалидную форму. Сравнивайте outputs и loss с эталоном через allclose с явно выбранным допуском. После смены backend проверяйте допустимую численную погрешность, а не требуйте безусловной битовой идентичности. Производительность измеряйте после компиляционного прогрева и завершения асинхронной работы. Не выбирайте лучший единичный замер: сообщайте распределение задержки и условия машины.
Диагностика расхождений
Если результат изменился, локализуйте слой: данные и preprocessing, key-поток, shapes/dtype, forward, loss, gradients, update, затем jit и backend. Сначала выполните функцию без jit на минимальном входе, затем верните transformations по одной. Сравните tree structure и каждый leaf. Не увеличивайте допуск и не применяйте nan_to_num, пока не выяснено, не скрывает ли это неверную форму, переполнение или повтор ключа.
Контроль после внедрения
В эксплуатации ограничьте допустимые batch shapes и объём входа, чтобы запрос не исчерпал память и не создал бесконечный поток перекомпиляций. Собирайте агрегированные метрики latency, compile count, долю невалидных входов и численные ошибки без записи исходных чувствительных данных. Сохраняйте предыдущий проверенный артефакт и код параметров для отката. Автоматическая замена весов требует той же независимой проверки, что и обычный выпуск.
Границы применимости
Пример на CPU не доказывает свойства GPU или TPU, а маленький массив не предсказывает производительность большой модели. JIT может ускорить повторяющиеся вычисления, но стоимость компиляции, передача данных и новые формы способны ухудшить end-to-end время. Autodiff возвращает производную написанной функции, а не подтверждает правильность математической постановки. Для критичных решений нужны предметные метрики, мониторинг и человеческий контроль.
Типичные ошибки и чего не делать
- Не используйте Python if с boolean tracer, зависящим от значения массива.
- Не возвращайте разные структуры или формы из ветвей.
- Не предполагайте, что нетрассируемая ветка может содержать произвольный Python-код.
Не публикуйте только удачный вывод без негативного теста и версий среды. Не используйте тестовую выборку для настройки модели. Не загружайте непроверенные параметры и не передавайте в диагностику персональные данные. Для важных решений предусмотрите безопасный отказ.
FAQ
Какой формы должен быть pred?
Это scalar boolean predicate.
Выполняются ли обе ветви?
Семантически выбирается одна, но обе ветви трассируются для построения программы.
Почему не Python if?
Runtime-значение tracer нельзя преобразовать в обычный Python bool под jit.
Официальные источники
Частые вопросы
Какой формы должен быть pred?
Это scalar boolean predicate.
Выполняются ли обе ветви?
Семантически выбирается одна, но обе ветви трассируются для построения программы.
Почему не Python if?
Runtime-значение tracer нельзя преобразовать в обычный Python bool под jit.
Читайте также
Как инициализировать двухслойную нейросеть в JAX
Двухслойный MLP без фреймворка: random.split, масштаб весов, ReLU, PyTree параметров, одиночный и пакетный входы.
Как найти NaN в JAX-нейросети через debug_nans
Локальная диагностика NaN в eager и jit: контекст jax.debug_nans, FloatingPointError, безопасная функция и запрет маскировки проблемы nan_to_num.
Как обновить веса нейросети через value_and_grad в JAX
Один проверяемый шаг gradient descent для PyTree параметров: MSE, value_and_grad, tree.map, формы градиентов и снижение loss.
Комментарии
Пока тихо. Скажите первое слово