Я занимаюсь рефакторингом проекта, написанного на Jax.
Есть функция, назовем ее foo, которая передается в jax.lax.scan . У него есть аргумент bar, который в настоящее время является частью переноса (т. е. первый аргумент, который представляет собой кортеж различных переменных, который передается следующему вызову). Я заметил, что аргумент bar не не изменяется за одно сканирование, т. е. функция распаковывает его из полученного переноса и упаковывает в возвращенный перенос без изменений. Я подумал, что лучше удалить его из переноса, но не мог понять, как это сделать.
Сначала я попробовал удалить его из переноса, добавив в качестве аргумента ключевого слова в функцию и изменив вызов сканирования на использование частичного(foo, bar=bar). Однако я заметил, что это медленно. Функция foo подвергается смещению, и я предполагаю, что способ, которым я добавил аргумент, заставляет его перемещаться каждый раз, а не только один раз.
Затем я попытался передать bar в аргумент xs, но у меня есть встроенные функции.ValueError: сканирование получено значение без ведущей оси для сканирования: 0, 0, 0, 0, 0, 0, 0, 0. >
Есть ли рекомендуемый способ сделать это?
Подробнее здесь: https://stackoverflow.com/questions/784 ... d-function
