Jax.lax.cond выполняет обе ветки вместо только ветки True ⇐ Python

[Anonymous]

Я обнаружил ошибку AttributeError: объект 'NoneType' не имеет атрибута в моем коде Python, использующем jax.lax.cond JAX. Вот упрощенная версия того, что я пытаюсь сделать:

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

from typing import NamedTuple, Optional
import jax

class Node(NamedTuple):
child: Optional['Node']
active: bool

node = Node(child=None, active=True)

child_active = jax.lax.cond(node.child is not None, lambda _: node.child.active, lambda _: False, None)

В этом коде node является экземпляром Node с дочерним элементом, установленным в значение None и активным, установленным в значение Верно. Согласно документации, jax.lax.cond должен выполнять только ту ветвь, которая имеет значение True. Однако я получаю ошибку AttributeError, указывающую, что ветвь False также выполняется.
Может ли кто-нибудь объяснить, почему эта ошибка возникает, когда jax.lax.cond следует оценивать только ветку True? Кроме того, каков правильный способ условного доступа к атрибутам необязательного дочернего узла, подобного этому, с помощью JAX?

Подробнее здесь: https://stackoverflow.com/questions/786 ... rue-branch