Flax中self.put_variable方法如何实现梯度穿透追踪

self.put_variable是Flax用于向模块外部传递内部可变状态的接口，不会在JAX自动微分计算图中建立赋值依赖：它既不会修改当前作用域已经获取到的参数引用，也不会让后续计算自动关联到你写入的新值。
你代码中W梯度全0本质和put_variable无关：你通过self.param拿到初始b之后，返回值jnp.sum(x+b)始终用的是这个初始b，你计算的(x@W+b).reshape(5,)仅被写入输出变量字典，完全没参与最终loss的计算，W不在反向传播路径上，梯度自然全0。

计算图内的参数更新直接用显式张量运算（推荐） #

JAX是纯函数式框架，不存在隐式原地赋值，所有需要梯度追踪的计算直接显式用JAX张量运算完成即可，不需要依赖put_variable做计算流内的变量传递。
如果需要把更新后的参数传出模块供外部使用（比如外循环参数更新），计算完新值后再调用put_variable写入即可——此时新值本身已经在计算路径上，梯度会正常回传。
修正后的可运行代码：

import jax 
from jax import numpy as jnp 
from jax import grad,random
import flax 
from flax import linen as nn 


class network(nn.Module):
    input_size : int 
    output_size : int 
    @nn.compact
    def __call__(self,x):
        W = self.param('W',nn.initializers.normal(),(self.input_size,self.output_size))
        b = self.param('b',nn.initializers.normal(),(self.output_size,))

        # 显式计算更新后的b，所有运算被JAX追踪，梯度自然连通
        new_b = (x@W+b).reshape(5,)
        # 仅用于把更新后的b传出模块，不承担计算图连接的作用
        self.put_variable("params","b",new_b)  
    
        # 后续计算使用更新后的new_b，保证W在计算路径上
        return jnp.sum(x+new_b)


if __name__ == "__main__":
    key = random.PRNGKey(0)
    key_x,key_param,key = random.split(key,3)
    x = random.normal(key_x,(1,5))

    module = network(5,5)
    param = module.init(key_param,x)
    loss, grads = grad(module.apply,has_aux=True)(param,x,mutable=["params"])
    print(grads)

运行后W会输出非零梯度，符合预期。

无梯度状态更新直接用put_variable #

如果你更新的是不需要梯度的状态（比如BatchNorm滑动均值、推理缓存等），直接通过self.variable定义变量、用put_variable更新即可，这类变量本身不需要梯度穿透，不会影响可训练参数的反向传播。

• 不要用put_variable实现计算图内的变量隐式替换，它仅负责跨模块边界传递状态，不会修改当前作用域的Python变量引用，也不会自动建立计算依赖
• 所有需要反传梯度的运算，必须保证最终输出和目标参数之间存在显式的张量计算依赖，JAX自动微分只会追踪实际参与运算的张量路径
• 开启mutable=["params"]回传更新参数时，回传的参数字典内的张量和计算图是连通的，只要内部计算显式使用了新值，梯度就不会中断

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