TensorFlow避免重复计算的程度及函数嵌套调用场景下的计算复用问题

首先来回答第一个问题：TensorFlow避免重复计算的能力，主要取决于你使用的是Eager Execution（即时执行，TF2.x默认）还是静态计算图模式（通过tf.function启用，TF1.x默认），两者的表现差异很大：

一、TensorFlow的通用重复计算避免逻辑 #

1. 静态计算图模式（tf.function装饰） #

在这种模式下，TensorFlow会先把你的代码转换成一个静态的计算图，然后再执行。它会自动检测图中输入完全相同、操作逻辑完全一致的子计算节点，将它们合并成一个，只执行一次计算，然后把结果复用给所有需要的地方。这种优化属于公共子表达式消除（Common Subexpression Elimination），是TensorFlow图优化的核心能力之一，能有效避免无意义的重复计算。

不过这种复用有前提：

• 两次调用的操作必须完全一致（包括参数、内部逻辑，甚至显式指定的操作名称也要一致）；
• 输入的张量必须是同一个（或者值完全相同且在图中被视为等价的输入）。如果输入有任何差异，哪怕是微小的，TensorFlow都会视为不同的计算分支，不会复用。

2. Eager Execution模式（默认） #

即时执行是逐行运行代码，默认情况下不会自动缓存任何计算结果。每次调用函数或者执行操作，都会重新跑一遍对应的逻辑，哪怕输入完全相同。如果你想在Eager模式下避免重复计算，需要手动实现缓存逻辑，比如用Python的functools.lru_cache（注意要处理张量的哈希问题，比如把张量转换成numpy数组作为键），或者使用TensorFlow提供的tf.cache（它主要在tf.function内生效）。

二、你的代码场景具体分析 #

来看你给出的代码片段：

import numpy as np
import tensorflow as tf
class MyClass:
    def all_eig(A):
        eigvals, eigvecs = tf.self_adjoint_eig(A)
        return eigvals
    def dominant_eig(A):
        eigvals = MyClass.all_eig(A)
        idx = tf.cast(tf.argmax(tf.abs(eigvals)), dtype=np.int32)
        return eigvals[idx]
def f(A):
    dom_eigval = MyClass.dominant_eig(A)
    eigvals = MyClass.all_eig(A)  # 这里的A和上面的是同一个输入
    # 后续逻辑...

我们分两种情况讨论：

1. 默认Eager模式下 #

当你调用f(A)时：

• 首先执行MyClass.dominant_eig(A)，里面会调用all_eig(A)，执行一次tf.self_adjoint_eig(A)计算特征值；
• 然后f函数里又直接调用MyClass.all_eig(A)，这时候会再次执行tf.self_adjoint_eig(A)，也就是重复计算了。因为Eager模式下没有自动缓存，每次函数调用都会重新执行内部的所有操作。

2. 使用tf.function装饰后的静态图模式 #

如果你把相关函数用@tf.function装饰，比如：

class MyClass:
    @tf.function
    def all_eig(A):
        eigvals, eigvecs = tf.self_adjoint_eig(A)
        return eigvals
    @tf.function
    def dominant_eig(A):
        eigvals = MyClass.all_eig(A)
        idx = tf.cast(tf.argmax(tf.abs(eigvals)), dtype=np.int32)
        return eigvals[idx]

@tf.function
def f(A):
    dom_eigval = MyClass.dominant_eig(A)
    eigvals = MyClass.all_eig(A)
    # 后续逻辑...

这时候TensorFlow在构建计算图时，会检测到两次all_eig(A)的输入是同一个张量A，操作逻辑完全一致，所以会将这两个调用合并成一个tf.self_adjoint_eig(A)的计算节点，只执行一次计算，然后把结果同时提供给dominant_eig和f中的all_eig调用，避免了重复计算。

不过要注意：如果A是每次调用f时都生成的新张量（哪怕值相同），只要在图中被视为相同的输入，就会复用；但如果A包含Python标量或者非TensorFlow类型的变量，可能会触发图的重新构建，影响优化效果。

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