torch.cond

torch.cond(pred, true_fn, false_fn, operands=())

有条件地应用 true_fn 或 false_fn。

警告

torch.cond 是 PyTorch 中的一个原型功能。目前它对输入和输出类型的支持有限，且不支持训练。敬请期待 PyTorch 未来版本中更稳定的实现。阅读更多关于功能分类的信息：https://pytorch.ac.cn/blog/pytorch-feature-classification-changes/#prototype

cond 是一个结构化控制流算子。也就是说，它类似于 Python 的 if 语句，但对 true_fn、false_fn 和 operands 有限制，这些限制使其可以使用 torch.compile 和 torch.export 进行捕获。

假设 cond 参数的约束条件满足，cond 等效于以下内容

def cond(pred, true_branch, false_branch, operands):
    if pred:
        return true_branch(*operands)
    else:
        return false_branch(*operands)

参数

• pred (Union[bool, torch.Tensor]) – 一个布尔表达式或一个只有一个元素的 tensor，指示应用哪个分支函数。

• true_fn (Callable) – 一个可调用函数 (a -> b)，位于正在被跟踪的作用域内。

• false_fn (Callable) – 一个可调用函数 (a -> b)，位于正在被跟踪的作用域内。真分支和假分支的输入和输出必须一致，这意味着输入必须相同，并且输出必须具有相同的类型和形状。

• operands (Tuple of possibly nested dict/list/tuple of torch.Tensor) – true/false 函数的输入元组。如果 true_fn/false_fn 不需要输入，它可以是空的。默认为 ()。

返回类型

Any

示例

def true_fn(x: torch.Tensor):
    return x.cos()
def false_fn(x: torch.Tensor):
    return x.sin()
return cond(x.shape[0] > 4, true_fn, false_fn, (x,))

限制条件

• 条件语句（即 pred）必须满足以下约束条件之一

  • 它是一个只有一个元素的 torch.Tensor，且数据类型为 torch.bool

  • 它是一个布尔表达式，例如 x.shape[0] > 10 或 x.dim() > 1 and x.shape[1] > 10

• 分支函数（即 true_fn/false_fn）必须满足以下所有约束条件

  • 函数签名必须与 operands 匹配。

  • 函数必须返回具有相同元数据（例如形状、数据类型等）的 tensor。

  • 函数不能对输入或全局变量进行原地修改。（注意：分支中允许使用原地 tensor 操作，例如用于中间结果的 add_）