如何在Flax中根据参数值选择执行不同函数？

遍历每个注意力头时，希望根据self.alpha的值选择执行f1或f2函数，仅执行其中一个，但当前代码运行时两个函数都会被执行。

原代码 #

def f1 (x):
    print('f1')
    return x/x.shape[2]
def f2 (x):
    print('f2')
    temp = nn.relu(x)
    return temp/(jnp.sum(temp,axis=-1,keepdims=True) + 1e-5)

def choose_attention(alpha, x):
    return jax.lax.cond(alpha[0, 0, 0,0],f2,f1,operand=x)

results = []
func = [f1,f2]
for i in range(self.alpha.shape[1]):
    print(i)
    alpha_i = self.alpha[:, i:i+1, :, :]
    x_i = attn_weights[:, i:i+1, :, :]
    result_i = jax.lax.switch(self.alpha[0,0,0,0].astype(int),func,x_i)
    results.append(result_i)

final_result = jnp.concatenate(results, axis=1)

打印输出 #

0
f1
f2
1
2
3
4
5
6
7
8
9
10
11

解决方案 #

核心原因 #

JAX是追踪式自动微分框架，追踪阶段会预求值所有分支代码，直接传入函数列表会导致所有函数都被执行。需要通过惰性包装确保仅选中分支执行。

方案1：用jax.lax.cond（更适合二选一场景） #

直接复用你定义的choose_attention函数，注意传入当前注意力头对应的alpha_i（而非固定值）：

results = []
for i in range(self.alpha.shape[1]):
    print(i)
    alpha_i = self.alpha[:, i:i+1, :, :]
    x_i = attn_weights[:, i:i+1, :, :]
    # 传入当前头的alpha值进行判断
    result_i = choose_attention(alpha_i, x_i)
    results.append(result_i)

final_result = jnp.concatenate(results, axis=1)

方案2：用jax.lax.switch（需包装函数为惰性执行） #

将函数用lambda包装，延迟执行时机：

results = []
# 用lambda包装函数，仅在被选中时执行
func = [lambda idx, x: f1(x), lambda idx, x: f2(x)]
for i in range(self.alpha.shape[1]):
    print(i)
    alpha_i = self.alpha[:, i:i+1, :, :]
    x_i = attn_weights[:, i:i+1, :, :]
    # 使用当前头对应的alpha值作为switch索引
    switch_idx = alpha_i[0,0,0,0].astype(int)
    result_i = jax.lax.switch(switch_idx, func, x_i)
    results.append(result_i)

final_result = jnp.concatenate(results, axis=1)

额外注意 #

原代码中switch使用固定的self.alpha[0,0,0,0]，会导致所有注意力头共用同一个判断条件，修改后改为当前循环的alpha_i对应值，确保每个头独立判断。

内容的提问来源于stack exchange，提问作者Naren Dhyani