Jax与Flax兼容问题：堆叠LSTM网络报错求助

错误根源 #

1. Jax与Flax变换冲突：在Flax的nn.Module内部直接使用jax.lax.scan会触发JaxTransformError——Flax的参数管理和自动微分依赖自身的变换系统，和原生Jax变换不兼容，必须使用Flax提供的flax.linen.scan。
2. 动态创建LSTMCell的问题：原代码在循环中每次创建nn.LSTMCell实例，会导致@nn.compact重复注册参数，引发参数命名冲突或初始化错误。
3. scan使用方式错误：替换flax.linen.scan后未正确配置参数（如参数共享、RNG处理），导致后续错误。

修正方案 #

关键修改点 #

• 用nn.scan替代jax.lax.scan，配置variable_broadcast='params'实现LSTM参数在序列步中的共享，split_rngs={'params': False}避免不必要的RNG分裂。
• 将LSTMCell的定义移到循环内但仅创建一次，确保参数只注册一次。
• 调整carry的更新逻辑，保证每层LSTM的状态正确传递。

修正后的完整代码 #

import jax
import jax.numpy as jnp
from flax import linen as nn
from typing import Sequence


class LSTMModel(nn.Module):
    lstm_hidden_size: int
    num_lstm_layers: int
    linear_layer_sizes: Sequence[int]
    mean_aggregation: bool

    def initialize_carry(self, batch_size, feature_size=1):
        """Initialize carry states with zeros for all LSTM layers."""
        return [
            (
                jnp.zeros((batch_size, self.lstm_hidden_size)),  # Hidden state (h)
                jnp.zeros((batch_size, self.lstm_hidden_size)),  # Cell state (c)
            )
            for _ in range(self.num_lstm_layers)
        ]

    @nn.compact
    def __call__(self, x, carry=None):
        if carry is None:
            raise ValueError("Carry must be initialized explicitly using `initialize_carry`.")

        # Expand 2D input to 3D (if necessary)
        if x.ndim == 2:
            x = jnp.expand_dims(x, axis=-1)

        # Process through LSTM layers
        for i in range(self.num_lstm_layers):
            # 定义当前层的LSTMCell，确保参数只注册一次
            lstm_cell = nn.LSTMCell(features=self.lstm_hidden_size, name=f'lstm_cell_{i}')
            
            # 用nn.scan包裹step_fn，配置参数共享与序列维度扫描
            @nn.scan(
                variable_broadcast='params',
                split_rngs={'params': False},
                in_axes=1,  # 对输入的第1维（序列维度）进行循环
                out_axes=1
            )
            def scan_step(carry, xt):
                new_carry, yt = lstm_cell(carry, xt)
                return new_carry, yt

            # 执行scan，更新当前层的carry和输出
            carry[i], outputs = scan_step(carry[i], x)
            x = outputs  # 将当前层输出作为下一层输入

        # 序列输出聚合
        if self.mean_aggregation:
            x = jnp.mean(x, axis=1)
        else:
            x = x[:, -1, :]

        # 全连接层处理
        for size in self.linear_layer_sizes:
            x = nn.Dense(features=size)(x)
            x = nn.elu(x)

        # 最终输出层
        x = nn.Dense(features=1)(x)
        return x


# 模型超参数
lstm_hidden_size = 64
num_lstm_layers = 2
linear_layer_sizes = [32, 16]
mean_aggregation = False

# 初始化模型
model = LSTMModel(
    lstm_hidden_size=lstm_hidden_size,
    num_lstm_layers=num_lstm_layers,
    linear_layer_sizes=linear_layer_sizes,
    mean_aggregation=mean_aggregation
)

# 生成测试输入
key = jax.random.PRNGKey(0)
dummy_input = jax.random.normal(key, (32, 10, 1))  # [batch_size, sequence_length, feature_size]

# 初始化隐藏状态
carry = model.initialize_carry(batch_size=dummy_input.shape[0])

# 初始化模型参数
params = model.init(key, dummy_input, carry)

# 前向传播
outputs = model.apply(params, dummy_input, carry)

print("Model output shape:", outputs.shape)  # 预期输出: (32, 1)

代码解释 #

1. nn.scan配置：variable_broadcast='params'指定LSTM的参数在整个序列的所有时间步中共享，避免重复初始化；in_axes=1表示对输入的序列维度（第1维）进行循环扫描。
2. LSTMCell定义：在每层循环中只创建一次LSTMCell实例，配合@nn.compact确保参数只注册一次，避免命名冲突。
3. 状态更新：每层scan执行后，更新当前层的carry状态，并将输出作为下一层的输入，保证多层LSTM的堆叠逻辑正确。

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