PyMC3上下文管理器存储模型参数的底层逻辑及最小示例咨询

你的猜测完全正确，这套自动注册变量的逻辑核心由两部分配合实现：

1. Model类的上下文管理器接口实现

• Model内部维护了一个线程独立的全局上下文栈，每次进入with块（即__enter__方法被调用）时，当前Model实例会被压入栈顶
• 退出with块（即__exit__方法被调用）时，当前实例会从栈中弹出，恢复上下文栈之前的状态

2. PyMC3分布类的初始化逻辑
   所有继承自PyMC3分布式基类的对象（比如pm.Normal、pm.HalfNormal）在初始化时，都会主动查询当前上下文栈的栈顶元素，如果存在有效Model实例，就自动把自己添加到该Model的变量存储列表中，同时完成计算图的关联注册。

以下示例完整复现该机制的核心逻辑，无额外依赖：

# 模拟全局上下文栈，PyMC3实际实现中会做线程隔离处理
context_stack = []

class Model:
    def __init__(self):
        self.vars = []  # 存储当前模型下所有注册的变量
    
    def __enter__(self):
        # 进入上下文时把当前模型压入栈顶
        context_stack.append(self)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        # 退出上下文时弹出当前模型，恢复之前的上下文状态
        if context_stack and context_stack[-1] is self:
            context_stack.pop()
        # 异常直接透传，不需要额外处理
        return False

# 模拟PyMC3的分布基类
class Distribution:
    def __init__(self, name, **kwargs):
        self.name = name
        self.params = kwargs
        # 初始化时自动注册到当前栈顶的Model实例
        if context_stack:
            current_model = context_stack[-1]
            current_model.vars.append(self)

# 模拟具体的Normal分布类
class Normal(Distribution):
    pass

# 测试逻辑，和PyMC3官方用法完全一致
if __name__ == "__main__":
    basic_model = Model()
    
    with basic_model:
        alpha = Normal("alpha", mu=0, sigma=10)
        beta = Normal("beta", mu=0, sigma=10, shape=2)
    
    # 打印验证变量已自动注册到模型
    print("模型中已存储的变量：")
    for var in basic_model.vars:
        print(f"- {var.name}, 参数：{var.params}")

注：PyMC3实际源码的实现比上述最小示例更复杂，会额外处理多线程上下文隔离、变量重名校验、计算图绑定、观测值自动关联等逻辑，但核心的上下文栈、自动注册逻辑和示例完全一致。

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