Flax升级后自定义RNNCell中initialize_carry兼容问题求助

以下是几种无需依赖setup方法、保留@nn.compact装饰器的可行方案：

方案1：通过carry状态判断延迟初始化子cell #

在__call__方法内实例化子RNNCell后，先检查传入的carry是否已经存在，仅在未初始化时调用子cell的initialize_carry方法，避免覆盖已有carry：

import flax.linen as nn
import jax.numpy as jnp

class CustomRNNCell(nn.Module):
  hidden_dim: int

  @nn.compact
  def __call__(self, carry, inputs):
    # 在compact模式下实例化子RNNCell
    lstm_cell = nn.LSTMCell(self.hidden_dim)
    
    # 仅当carry未传入时执行初始化
    if carry is None:
      batch_size = inputs.shape[0]
      carry = lstm_cell.initialize_carry(self.make_rng('params'), (batch_size,))
    
    # 执行子cell前向计算，复用传入的carry
    new_carry, output = lstm_cell(carry, inputs)
    return new_carry, output

方案2：用自定义变量标记初始化状态 #

通过Flax的variable创建一个非参数变量，记录子cell的carry是否已完成初始化，避免重复初始化覆盖carry：

import flax.linen as nn
import jax.numpy as jnp

class CustomRNNCell(nn.Module):
  hidden_dim: int

  @nn.compact
  def __call__(self, carry, inputs):
    lstm_cell = nn.LSTMCell(self.hidden_dim)
    # 创建标记变量，记录是否已初始化carry
    is_initialized = self.variable('carry_init', 'is_initialized', lambda: False)
    
    # 仅在carry未传入且未初始化过的情况下执行初始化
    if carry is None and not is_initialized.value:
      batch_size = inputs.shape[0]
      carry = lstm_cell.initialize_carry(self.make_rng('params'), (batch_size,))
      is_initialized.value = True
    
    new_carry, output = lstm_cell(carry, inputs)
    return new_carry, output

方案3：拆分carry结构兼容自定义逻辑 #

如果自定义cell需要额外的carry状态，可将传入的carry设计为包含子cell carry的复合结构，明确区分自定义状态与子cell状态，避免初始化冲突：

import flax.linen as nn
import jax.numpy as jnp

class CustomRNNCell(nn.Module):
  hidden_dim: int
  custom_carry_dim: int

  @nn.compact
  def __call__(self, carry, inputs):
    lstm_cell = nn.LSTMCell(self.hidden_dim)
    
    # 拆分复合carry：(自定义状态, 子cell状态)
    if carry is None:
      batch_size = inputs.shape[0]
      # 分别初始化自定义carry和子cell carry
      custom_carry = jnp.zeros((batch_size, self.custom_carry_dim))
      lstm_carry = lstm_cell.initialize_carry(self.make_rng('params'), (batch_size,))
      carry = (custom_carry, lstm_carry)
    
    custom_carry, lstm_carry = carry
    # 执行子cell计算，复用传入的子cell carry
    new_lstm_carry, lstm_output = lstm_cell(lstm_carry, inputs)
    # 自定义carry的更新逻辑
    new_custom_carry = custom_carry + jnp.mean(lstm_output, axis=-1, keepdims=True)
    
    return (new_custom_carry, new_lstm_carry), lstm_output

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