jax

37bb6f09 - partial evaluation rule for single-operand cond primitive