В Jax будет ли JIT-компилятор разворачивать циклы внутри оператора if, который не выполняется? Например,
def f(x):
if False:
for i in range(100):
do something
rest of function...
jax.jit(f)
Развертывание цикла было бы нежелательно, поскольку это может значительно увеличить время компиляции.
Из некоторых поисков в Интернете и вопросов ИИ кажется, что ответ должен быть отрицательным. То есть компилятор не увидит код в этом операторе if, поэтому цикл не развернется. Однако существует противоречивая информация, предполагающая, что на всякий случай следует использовать условные выражения jax. Каков окончательный ответ?
Приведенный выше пример совершенно надуман, чтобы проиллюстрировать эту идею. В общем, я ищу ответ для случая, когда условие оператора if можно полностью определить на основе статических аргументов функции. Например,
def f(x, do_the_loop):
if do_the_loop:
for i in range(100):
do something
rest of function...
jax.jit(f, static_argnums=1)
f(x, False)
