Зачем нужен jax.vmap и когда он лучше бродкастинга

Разбираем, когда в JAX использовать jax.vmap вместо бродкастинга: векторизация без Python‑циклов, пример с jnp.histogram, работа с PyTrees и производительность.

JAX и векторизация поначалу могут показаться запутанными. Многие операции с массивами уже принимают батчи за счёт бродкастинга, так почему вообще существует jax.vmap? Короткий ответ: не каждая операция нативно векторизована, и нередко явная векторизация получается понятнее, безопаснее и быстрее, чем манипуляции с размерностями или циклы на Python.

Когда «уже векторизовано» — этого недостаточно

Бродкастинг и батч-оси широко используются в коде для глубокого обучения и часто покрывают типовые случаи. Но есть операции в JAX, которые не работают нативно по батч-оси. Для них либо пишут явный цикл, либо применяют jax.vmap, чтобы получить батч-версию функции. Такой подход сохраняет читаемость и обычно быстрее, чем циклы на стороне Python.

Демонстрация проблемы

Возьмём jnp.histogram. Она не принимает батч-ось нативно. Прямолинейный вариант — пройтись циклом в Python и вручную собрать результаты в стек.

import jax
import jax.numpy as jnp
# Вспомогательная функция для гистограммы одного образца
def single_hist(vec, nbins, span):
    # jnp.histogram возвращает (hist, bin_edges); берём только гистограмму
    counts, _ = jnp.histogram(vec, bins=nbins, range=span)
    return counts
# Игрушечный батч образцов (2 образца, по 3 значения каждый)
batch_samples = jnp.array([[0.0, 1.0, 2.0],
                           [2.0, 3.0, 1.0]])
num_bins = 4
value_range = (0.0, 4.0)
# Батчинг на стороне Python (цикл + stack)
looped = [single_hist(row, num_bins, value_range) for row in batch_samples]
batched_counts = jnp.stack(looped, axis=0)

Это работает, но добавляет цикл на Python и лишние телодвижения со стекированием. К тому же легко ошибиться, когда вручную добавляешь или переставляешь оси ради имитации батча.

В чём корень проблемы?

Путаница возникает из‑за того, что «векторизовано» означает разные вещи: одни примитивы изначально умеют обрабатывать батч-ось, другие по умолчанию скалярные или «одиночного образца» и требуют батчинга. jnp.histogram и jnp.bincount относятся ко второй группе. В таких случаях jax.vmap даёт аккуратный способ выразить «применять эту функцию независимо к каждому элементу батча», не меняя её семантику для одного образца и не извращаясь с формами.

Есть и стилистический аспект. Иногда разработчики добавляют временную дополнительную ось, чтобы избежать цикла, а затем её сворачивают. Работает, но намерение обычно яснее, когда это прямо обозначено через jax.vmap.

Решение с jax.vmap

jax.vmap преобразует функцию для одного образца в батч-функцию по выбранной оси. Он работает с PyTrees, что позволяет библиотекам скрывать обработку батча целиком. Для независимых операций над образцами это естественный выбор.

import jax
import jax.numpy as jnp
# Та же функция для одного образца
def single_hist(vec, nbins, span):
    counts, _ = jnp.histogram(vec, bins=nbins, range=span)
    return counts
batch_samples = jnp.array([[0.0, 1.0, 2.0],
                           [2.0, 3.0, 1.0]])
num_bins = 4
value_range = (0.0, 4.0)
# Векторизованный вариант: применяем по ведущей оси batch_samples
vmapped_hist = jax.vmap(lambda v: single_hist(v, num_bins, value_range))
counts_batched = vmapped_hist(batch_samples)

Вариант с vmap отражает ту же логику, что и реализация с циклом, но без явного цикла и ручного стекирования. Это удобнее и помогает избежать циклов, улучшая читаемость и, как правило, производительность.

Где vmap особенно полезен сверх базового

vmap работает с PyTrees, поэтому можно обрабатывать целые структуры параметров, не пиша код для управления осями. Некоторые библиотеки, например equinox, опираются на такую конвенцию и поощряют vmapping по всему дереву параметров. Это снимает необходимость вручную протаскивать батч-оси через код модели. Подход предполагает независимость между образцами и не подходит для операций, которые принципиально смешивают информацию между ними, например батч-нормализации.

В других ситуациях может возникнуть соблазн добавить временную ось, заставить операцию отработать за счёт бродкастинга, а затем убрать лишнее измерение. Часто vmap напрямую выражает этот замысел. Для интуиции: представьте применение convolution2d с разными ядрами для каждого образца. Один путь — сложить ядра в стек, реплицировать и сложить каналы, а затем выполнить единственную свёртку по расширенной оси. Другой — написать свёртку для одного образца и провезти её через vmap по оси ядер или по оси образцов. Оба способа могут работать; vmap просто явно фиксирует независимость по образцам.

Почему это важно

Стратегия батчинга — не только вопрос стиля; она влияет на корректность, читаемость и производительность. Понимание, когда операция не векторизована нативно, помогает избежать тихих ошибок форм и ненароком добавленных Python-циклов, которые урезают пропускную способность. Применяя vmap там, где это уместно, вы сохраняете логику «на один образец», легко комбинируете её с PyTrees и следуете соглашениям библиотек, полагающихся на vmap.

Практические рекомендации

В повседневном коде бродкастинг и батч-оси — хороший базовый выбор. Тянитесь к jax.vmap, когда встречаете функции без нативной векторизации, когда дизайн библиотеки поощряет vmapping по PyTrees или когда нужно векторизовать по нетипичным осям. Если ловите себя на добавлении разовых измерений и последующих редукций, чтобы «обойтись без цикла», подумайте, не сделает ли vmap намерение прозрачнее, а код — удобнее в сопровождении.

Универсального правила нет, и играет роль личное предпочтение. Важно видеть границу: используйте нативный батчинг там, где он есть и уместен, а vmap — чтобы выражать независимые вычисления по образцам, когда его нет. Такая дисциплина ведёт к более простому и надёжному коду на JAX.