Ошибка типа: объект «Добавить» не является итеративным при переборе выражений SymPy в scqubits с помощью JAX custom_vjp ⇐ Python

[Anonymous]

Я пытаюсь изменить пакет Python scqubits, чтобы использовать JAX custom_vjp для дифференцируемого программирования. При этом я столкнулся со следующей ошибкой:
TypeError: объект «Добавить» не повторяется
Вот минимальный пример, воспроизводящий проблему:

Код: Выделить всё

import sympy as sm
x, y, z = sm.symbols('x y z')

expr = x**2 + 2*y + z

try:
for term in expr:
print(term)
except TypeError as e:
print(f"Error: {e}")

Этот код выводит:
Ошибка: объект «Добавить» не повторяется
В базе кода scqubits, особенно в scqubits/core/ Circuit_routines.py есть функция, которая перебирает выражение SymPy:

Код: Выделить всё

def _constants_in_subsys(self, H_sys: sm.Expr, constants_expr: sm.Expr) -> sm.Expr:
"""
Returns an expression of constants that belong to the subsystem with the
Hamiltonian H_sys

Parameters
----------
H_sys:
Subsystem Hamiltonian

Returns
-------
Expression of constants belonging to the subsystem
"""
constant_expr = 0
subsys_free_symbols = set(H_sys.free_symbols)
constant_terms = constants_expr.copy()
for term in constant_terms:
if set(term.free_symbols) & subsys_free_symbols == set(term.free_symbols):
constant_expr += term
return constant_expr

Когда я запускаю эту функцию с помощью custom_vjp JAX, она вызывает ту же ошибку TypeError.
Мои вопросы:
Неправильно ли выполнять итерацию по выражению SymPy напрямую (как в случае с термином в expr)?
Почему возникает ошибка TypeError, говорящая, что объект «Добавить» не повторяется?
Как я могу изменить код так, чтобы он был корректным перебирать условия выражения SymPy, особенно в контексте использования JAX?
Дополнительная информация:
Версия SymPy: (1.9)
Версия JAX: (0.2) .25)
Версия scqubits: (если применимо)
Я пробовал заменить итерацию на expr.args, вот так:

Код: Выделить всё

for term in expr.args:
print(term)

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

Подробнее здесь: https://stackoverflow.com/questions/789 ... essions-in