Как использовать jax.custom_vjp с функциями, принимающими не-JAX типы (например, выражения SymPy) в качестве входных данных?

Содержание

Вопрос или проблема
Ответ или решение
Проблема
Решение
Объяснение кода

Вопрос или проблема

Я пытаюсь использовать custom_vjp из JAX для определения пользовательских вычислений градиентов для функции, которая принимает выражение SymPy в качестве входного параметра. Однако я сталкиваюсь с ошибками, потому что JAX не поддерживает типы, не относящиеся к JAX, в качестве входных параметров для функций, которые трансформируются (например, с помощью grad, jit или custom_vjp). В последнее время я модифицирую код в ScQubits, чтобы добавить новый бэкенд jax для повышения эффективности, и затем сталкиваюсь с некоторыми проблемами с jax и sympy.

Вот минимальный пример того, что я пытаюсь сделать:

import jax
import sympy as sm

# Определите символы и выражение
x, y, z = sm.symbols('x y z')
expr = x**2 + 2*y + z

# Попытка итерации по expr (это вызовет ошибку)
try:
    for term in expr:
        print(term)
except TypeError as e:
    print(f"Ошибка: {e}")

# Определите функцию, которая принимает выражение SymPy и значение
def sympy_function(expr, x_value):
    x = sm.Symbol('x')
    result = expr.subs(x, x_value)
    return float(result)

# Попытка применить custom_vjp
sympy_function = jax.custom_vjp(sympy_function)

def sympy_function_fwd(expr, x_value):
    y = sympy_function(expr, x_value)
    return y, (expr, x_value)

def sympy_function_bwd(residual, grad_y):
    expr, x_value = residual
    x = sm.Symbol('x')
    derivative_expr = sm.diff(expr, x)
    grad_x_value = float(derivative_expr.subs(x, x_value))
    grad_expr = None
    return grad_expr, grad_y * grad_x_value

sympy_function.defvjp(sympy_function_fwd, sympy_function_bwd)

# Тестирование функции
x = sm.Symbol('x')
expr = x**2 + 3*x + 2
x_value = 1.0

# Это вызовет ошибку
y = sympy_function(expr, x_value)

Когда я запускаю этот код, я получаю ошибку подобного рода:

TypeError: Значение x**2 + 3*x + 2 с типом <class 'sympy.core.add.Add'> не является допустимым типом JAX

Как я могу использовать jax.custom_vjp с функциями, которые принимают нетипизированные типы, такие как выражения SymPy в качестве входных параметров? Есть ли способ обойти это ограничение или сделать так, чтобы JAX принимал такие функции?

Ответ или решение

Конечно, давайте разберем, как можно обойти ограничения JAX в отношении использования нестандартных типов, таких как выражения SymPy, в функции с пользовательским градиентом на основе jax.custom_vjp.

Проблема

Как вы уже заметили, попытка использовать выражения SymPy в функции, обернутой в JAX, вызывает ошибку, поскольку JAX ожидает работы со своими собственными типами данных и массивами. Поэтому, чтобы использовать jax.custom_vjp с SymPy, нужно будет преобразовать выражения в формат, который JAX сможет обрабатывать.

Решение

Конвертация SymPy в JAX: Передаем числовые значения, а для расчета производной используем JAX для обработки градиентов. Для этого мы будем использовать функции JAX для вычислений и преобразовывать выражения SymPy в функции Python.
Оборачивание SymPy: Мы можем преобразовать SymPy-выражения в функции, которые могут принимать значения и возвращать результаты, которые JAX сможет обрабатывать.

Вот исправленный код:

import jax
import jax.numpy as jnp
import sympy as sm

# Определяем символы для выражения
x, y, z = sm.symbols('x y z')
expr = x**2 + 2*y + z

# Преобразуем SymPy-выражение в функцию
def sympy_to_function(expr):
    var = sm.symbols('x')  # Замените на нужные переменные
    expr_lambdified = sm.lambdify(var, expr, modules='numpy')
    return expr_lambdified

# Преобразование выражения в функцию
sympy_function = sympy_to_function(expr)

# Определяем пользовательские функции для JAX
@jax.custom_vjp
def jax_sympy_function(x_value):
    return sympy_function(x_value)

# Прямой проход
def jax_sympy_function_fwd(x_value):
    y = jax_sympy_function(x_value)
    return y, x_value  # Возвращаем значение и входное значение для обратного прохода

# Обратный проход
def jax_sympy_function_bwd(residual, grad_y):
    x_value = residual
    x_sympy = sm.symbols('x')
    derivative_expr = sm.diff(expr, x_sympy)
    grad_x_value = float(derivative_expr.subs(x_sympy, x_value))
    return grad_y * grad_x_value  # Возвращаем градиент

# Связываем пользовательские функции с jax_sympy_function
jax_sympy_function.defvjp(jax_sympy_function_fwd, jax_sympy_function_bwd)

# Тестируем функцию
x_value = 1.0
y = jax_sympy_function(x_value)

print("Результат:", y)

# Проверка градиента
grad = jax.grad(jax_sympy_function)(x_value)
print("Градиент:", grad)

Объяснение кода

sympy_to_function(expr): Эта функция преобразует SymPy выражение в обычную функцию с использованием lambdify. Теперь мы можем вызывать эту функцию с числовыми значениями.
jax_sympy_function: Эта функция оборачивает преобразованную функцию, делая её совместимой с JAX.
jax_sympy_function_fwd и jax_sympy_function_bwd: Эти функции реализуют выполнение прямого и обратного прохода, где мы вычисляем значение и производную.

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