Как избежать TracerArrayConversionError в JAX jit при работе с состоянием

Почему изменение атрибутов в jitted‑методах JAX ведет к TracerArrayConversionError, когда значения материализуются, и как безопасно реорганизовать код.

JAX, jit и состояние объектов не всегда ладят. Распространённая ловушка — попытка менять атрибуты внутри jitted-метода, а затем обращаться к ним как к обычным массивам NumPy. В итоге вы ловите TracerArrayConversionError и начинаете гадать, когда же значения становятся конкретными и как организовать код, чтобы без неожиданностей.

Постановка задачи

Проблема проявляется, когда jitted‑функция записывает промежуточное значение в атрибут объекта, а затем к этому атрибуту обращаются за пределами jitted‑контекста. Пример ниже воспроизводит поведение.

def host_bootstrap():
    print("This code runs once at the beginning of the program.")
    return jnp.array([1.0, 2.0, 3.0])

class Box:

    @partial(jax.jit, static_argnums=0)
    def apply_jit(self, arr):
        print("This code runs once during the first trace.")
        out = arr * 2
        self.cache = out
        return out

# Выполнение программы
payload = host_bootstrap()
box = Box()
out1 = box.apply_jit(payload)  # Трассировка происходит здесь

np.array(out1)
np.array(box.cache)

Вызов np.array для результата метода срабатывает, но тот же приём для атрибута падает с:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[3]

Что на самом деле не так

Преобразования JAX, такие как jit, рассчитаны на чистые функции. Чистая функция не изменяет входы и не замыкает скрытое состояние; её результат зависит только от аргументов. Присваивание self.cache внутри jitted‑метода фактически мутирует вход функции и нарушает контракт чистоты. Во время трассировки/компиляции jit значения вроде out — это трассеры, а не конкретные массивы. Возврат значения из jitted‑функции задаёт JAX чёткую границу, где его можно материализовать; простое присваивание того же значения атрибуту — нет.

Когда значение в атрибуте станет конкретным?

Только когда оно возвращается из JIT‑скомпилированной функции. Если возврата нет, на этом этапе исполнения это остаётся трассером; попытка обращаться к нему как к массиву NumPy приводит к ошибке конверсии. Код, который меняет состояние внутри jit, ведёт себя неопределённо.

Есть и другой тревожный момент: помечать self как static и одновременно его менять. Так делать прямо не рекомендуется. См. документацию JAX о чистых функциях и раздел FAQ о том, как использовать jit в методах, — там разбираются типичные ловушки и предложены корректные подходы: Pure Functions и How to use jit with methods.

Как исправить дизайн

Если значение требуется за пределами jitted‑функции, верните его и обновите состояние объекта на стороне Python. Так вы сохраняете чистоту и даёте JAX понятную точку материализации результата. И, разумеется, не помечайте self как static, если планируете его изменять или от него зависеть.

def host_bootstrap():
    print("This code runs once at the beginning of the program.")
    return jnp.array([1.0, 2.0, 3.0])

class Box:

    @jax.jit
    def apply_jit(self, arr):
        print("This code runs once during the first trace.")
        out = arr * 2
        return out

# Выполнение программы
payload = host_bootstrap()
box = Box()
out1 = box.apply_jit(payload)

# Обновляем состояние объекта вне jit
box.cache = out1

# Обе конверсии теперь проходят нормально
np.array(out1)
np.array(box.cache)

Есть и другой рабочий вариант: сделать объект pytree, не помечать его static и возвращать из функции обновлённый экземпляр. Это поддерживает идею явной передачи состояния вместо скрытых побочных эффектов.

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

Модель компиляции и преобразований JAX исходит из чистых функций. Опора на побочные эффекты во время трассировки рушит эти предпосылки и порождает хрупкое, неопределённое поведение. Возврат значений задаёт понятную границу материализации и делает преобразования предсказуемыми и сочетаемыми.

Выводы

Нужно значение вне jitted‑метода — верните его. Не меняйте атрибуты внутри jit и не помечайте self как static, если он изменяется. Если состояние должно переходить между вызовами, передавайте его явно — или оформите класс как pytree и возвращайте обновлённое состояние. Эти приёмы соответствуют модели исполнения JAX и убирают сюрпризы вроде TracerArrayConversionError.