如何将多个Python函数融合为性能与手写硬编码相当的函数？

方案：动态生成表达式代码 #

原实现性能差的核心原因是循环调用每个运算函数存在多次函数调用开销，和手写的合并表达式相比多了N次函数调用的成本，我们可以通过预定义每个运算的表达式模板，动态拼接生成合并后的函数代码的方式实现和手写fastF完全一致的性能。

具体实现步骤 #

1. 给所有你预定义的基础运算添加表达式模板属性，用来描述该运算的表达式格式：

# 定义基础运算并绑定表达式模板
sq = lambda x: x**2
sq._expr_template = "({} ** 2)"

add = lambda x: x + 3
add._expr_template = "({} + 3)"

mul = lambda x: x * 5
mul._expr_template = "({} * 5)"

# 其他魔术方法对应的操作同理，比如矩阵乘法可以定义为：
# matmul_A = lambda x: x @ A
# matmul_A._expr_template = "({} @ A)"

2. 实现generateF函数，根据传入的操作顺序拼接表达式，动态生成目标函数：

def generateF(ops):
    # 从初始变量x开始拼接表达式
    current_expr = "x"
    for op in ops:
        current_expr = op._expr_template.format(current_expr)
    # 动态生成函数，和手写的fastF代码完全一致
    return eval(f"lambda x: {current_expr}")

效果验证 #

ops = [add, sq, mul, sq]
f = generateF(ops)
print(f(3)) # 输出32400，和fastF结果一致

生成的f本质上和你手写的fastF = lambda x: (5*(x+3)**2)**2完全相同，运行性能没有任何差异。

方案优势 #

• 无Python版本兼容问题，所有版本Python都支持eval语法
• 完全满足约束要求：不需要C扩展，支持PyTorch张量等任意类型输入，所有运算逻辑和你预定义的操作完全一致
• 安全性有保障：所有表达式模板都是你作为库开发者预定义的，用户仅能选择操作顺序，不存在代码注入风险
• 适配所有数值魔术方法：不管是加减乘除、移位、矩阵乘法，都可以提前定义对应的表达式模板即可

内容的提问来源于stack exchange，提问作者157 239n