espnet2.asr_transducer.decoder.modules.rwkv.feed_forward.FeedForward

source

class espnet2.asr_transducer.decoder.modules.rwkv.feed_forward.FeedForward(size: int, hidden_size: int, block_id: int, num_blocks: int)

Bases: Module

FeedForward module definition.

• Parameters:
  • size – Input/Output size.
  • hidden_size – Hidden size.
  • block_id – Block index.
  • num_blocks – Number of blocks in the architecture.

Construct a FeedForward object.

forward(x: Tensor, state: List[Tensor] | None = None) → Tuple[Tensor, List[Tensor] | None]

Compute channel mixing.

• Parameters:
  • x – FeedForward input sequences. (B, U, size)
  • state – Decoder hidden state. [5 x (B, 1, size, N)]
• Returns: FeedForward output sequences. (B, U, size) state: Decoder hidden state. [5 x (B, 1, size, N)]
• Return type: x

reset_parameters(size: int, block_id: int, num_blocks: int) → None

Reset module parameters.

• Parameters:
  • size – Block size.
  • block_id – Block index.
  • num_blocks – Number of blocks in the architecture.