为何Flax的Dropout实现采用jax.lax.select而非jax.numpy.where或乘法？

我查阅了Flax的Dropout实现，核心代码如下：

def __call__(self, inputs, deterministic: Optional[bool] = None):
    """Applies a random dropout mask to the input.

    Args:
      inputs: the inputs that should be randomly masked.
      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
        masked, whereas if true, no mask is applied and the inputs are returned
        as is.

    Returns:
      The masked inputs reweighted to preserve mean.
    """
    deterministic = merge_param(
        'deterministic', self.deterministic, deterministic)

    if (self.rate == 0.) or deterministic:
      return inputs

    # Prevent gradient NaNs in 1.0 edge-case.
    if self.rate == 1.0:
      return jnp.zeros_like(inputs)

    keep_prob = 1. - self.rate
    rng = self.make_rng(self.rng_collection)
    broadcast_shape = list(inputs.shape)
    for dim in self.broadcast_dims:
      broadcast_shape[dim] = 1
    mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
    mask = jnp.broadcast_to(mask, inputs.shape)
    return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))

我特别关注最后一行的lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))，想知道为什么要使用jax.lax.select，而不是以下两种更直观的写法：

写法一：

return jnp.where(mask, inputs / keep_prob, 0)

写法二：

return mask * inputs / keep_prob

为什么选择lax.select而非另外两种写法？ #

1. 和jnp.where的区别：类型一致性与广播控制

jnp.where是lax.select的高层封装，但传入标量0作为第三个参数时，JAX会自动将其广播为与inputs同形状的数组，同时可能触发隐式类型转换——比如如果inputs是bfloat16类型，标量0默认是float32，转换过程会带来额外开销，甚至可能引入精度损失。而jnp.zeros_like(inputs)会严格生成与输入同形状、同类型的零数组，完全避免了这个问题。

另外，lax.select作为底层API，行为更可控，不会有高层封装带来的额外隐式操作，这对追求稳定性和可预测性的框架核心代码来说很重要。

2. 和mask * inputs / keep_prob的区别：数值稳定性与计算效率

首先，mask是布尔数组，在乘法运算中会被自动转换为0./1.的浮点数组。当keep_prob很小时（比如dropout rate接近1），inputs / keep_prob会得到极大的数值，此时乘以0.可能会引入数值精度问题（比如极大值乘0可能得到非零的极小值，而非严格的0）。而lax.select会直接在mask为False时返回预先生成的零数组，完全避免这种数值异常。

其次，从自动微分的角度看，mask * inputs / keep_prob的梯度计算会涉及所有元素，即使mask为False的位置；而lax.select会在反向传播时跳过mask为False的分支，减少不必要的梯度计算开销，尤其是在大张量上表现更明显。

3. 框架代码的一致性与可维护性

Flax作为JAX生态的框架，倾向于直接使用底层laxAPI来保持代码的一致性，避免依赖高层API的潜在行为变化。这种写法也能让框架开发者更清晰地控制计算流程，便于后续的优化和维护。

内容的提问来源于stack exchange，提问作者Jay Mody