model

class MatrixGame3TimeImageEmbedding(nn.Module):
def __init__(
    self,
    dim: int,
    time_freq_dim: int,
):
    super().__init__()

    self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
    self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu")

def forward(
    self,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None,
    timestep_seq_len: int | None = None,
):
    temb = self.time_embedder(timestep, timestep_seq_len)
    timestep_proj = self.time_modulation(temb)

    if encoder_hidden_states is None:
        batch_size = temb.shape[0] if temb.dim() > 1 else timestep.shape[0]
        encoder_hidden_states = torch.zeros(
            (batch_size, 0, temb.shape[-1]),
            device=temb.device,
            dtype=temb.dtype,
        )

    return temb, timestep_proj, encoder_hidden_states</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3TimeImageEmbedding.forward
forward(timestep: torch.Tensor, encoder_hidden_states: torch.Tensor | None, timestep_seq_len: int | None = None)
Source
def forward(
    self,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None,
    timestep_seq_len: int | None = None,
):
    temb = self.time_embedder(timestep, timestep_seq_len)
    timestep_proj = self.time_modulation(temb)
if encoder_hidden_states is None:
    batch_size = temb.shape&#91;0&#93; if temb.dim() &gt; 1 else timestep.shape&#91;0&#93;
    encoder_hidden_states = torch.zeros(
        (batch_size, 0, temb.shape&#91;-1&#93;),
        device=temb.device,
        dtype=temb.dtype,
    )

return temb, timestep_proj, encoder_hidden_states</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention
MatrixGame3CrossAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, parallel_attention = False, quant_config: QuantizationConfig | None = None, prefix: str = '')
Source
class MatrixGame3CrossAttention(WanSelfAttention):
def forward(self, x, context, context_lens=None):
    r&quot;&quot;&quot;
    Args:
        x(Tensor): Shape &#91;B, L1, C&#93;
        context(Tensor): Shape &#91;B, L2, C&#93; - typically 257 image tokens
        context_lens(Tensor): Shape &#91;B&#93;
    &quot;&quot;&quot;
    b, n, d = x.size(0), self.num_heads, self.head_dim

    q_input = x.to(self.to_q.weight.dtype)
    q = self.norm_q(self.to_q(q_input)&#91;0&#93;).view(b, -1, n, d)

    context_input = context.to(self.to_k.weight.dtype)
    k = self.norm_k(self.to_k(context_input)&#91;0&#93;).view(b, -1, n, d)
    v = self.to_v(context_input)&#91;0&#93;.view(b, -1, n, d)

    x = self.attn(q, k, v)

    x = x.flatten(2)
    x, _ = self.to_out(x)
    return x</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention.forward
forward(x, context, context_lens = None)
Parameters

x: Tensor
Shape [B, L1, C]
context: Tensor
Shape [B, L2, C] - typically 257 image tokens
context_lens: Tensor
Shape [B]

Source
def forward(self, x, context, context_lens=None):
    r"""
    Args:
        x(Tensor): Shape [B, L1, C]
        context(Tensor): Shape [B, L2, C] - typically 257 image tokens
        context_lens(Tensor): Shape [B]
    """
    b, n, d = x.size(0), self.num_heads, self.head_dim
q_input = x.to(self.to_q.weight.dtype)
q = self.norm_q(self.to_q(q_input)&#91;0&#93;).view(b, -1, n, d)

context_input = context.to(self.to_k.weight.dtype)
k = self.norm_k(self.to_k(context_input)&#91;0&#93;).view(b, -1, n, d)
v = self.to_v(context_input)&#91;0&#93;.view(b, -1, n, d)

x = self.attn(q, k, v)

x = x.flatten(2)
x, _ = self.to_out(x)
return x</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3TransformerBlock
MatrixGame3TransformerBlock(dim: int, ffn_dim: int, num_heads: int, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', action_config: dict | None = None, block_id: int | None = None, use_memory: bool = False)
Source
class MatrixGame3TransformerBlock(nn.Module):
def __init__(
    self,
    dim: int,
    ffn_dim: int,
    num_heads: int,
    qk_norm: str = &quot;rms_norm_across_heads&quot;,
    cross_attn_norm: bool = False,
    eps: float = 1e-6,
    supported_attention_backends: tuple&#91;AttentionBackendEnum, ...&#93;
    | None = None,
    prefix: str = &quot;&quot;,
    action_config: dict | None = None,
    block_id: int | None = None,
    use_memory: bool = False,
):
    super().__init__()
    action_config = action_config or &#123;&#125;
    self.use_memory = use_memory

    # 1. Self-attention
    self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
    self.to_q = ReplicatedLinear(dim, dim, bias=True)
    self.to_k = ReplicatedLinear(dim, dim, bias=True)
    self.to_v = ReplicatedLinear(dim, dim, bias=True)

    self.to_out = ReplicatedLinear(dim, dim, bias=True)
    self.attn1 = DistributedAttention(
        num_heads=num_heads,
        head_size=dim // num_heads,
        causal=False,
        supported_attention_backends=supported_attention_backends,
        prefix=f&quot;&#123;prefix&#125;.attn1&quot;,
    )
    self.hidden_dim = dim
    self.num_attention_heads = num_heads
    dim_head = dim // num_heads
    if qk_norm == &quot;rms_norm&quot;:
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)
    elif qk_norm == &quot;rms_norm_across_heads&quot;:
        # LTX applies qk norm across all heads
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
    else:
        print(&quot;QK Norm type not supported&quot;)
        raise Exception
    assert cross_attn_norm is True
    self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
        dim,
        norm_type=&quot;layer&quot;,
        eps=eps,
        elementwise_affine=True,
        dtype=torch.float32,
        compute_dtype=torch.float32,
    )

    # 2. Cross-attention
    self.attn2 = MatrixGame3CrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
    self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
        dim,
        norm_type=&quot;layer&quot;,
        eps=eps,
        elementwise_affine=False,
        dtype=torch.float32,
        compute_dtype=torch.float32,
    )

    # 2.1. Action Module Integration
    enabled_blocks = set(action_config.get(&quot;blocks&quot;, &#91;&#93;)) if action_config else set()
    self.use_action_module = len(action_config) &gt; 0 and (block_id is None or block_id in enabled_blocks)
    if self.use_action_module:
        self.action_model = MatrixGame3ActionModule(**action_config)
    else:
        self.action_model = None

    # 3. Feed-forward
    self.ffn = MLP(dim, ffn_dim, act_type=&quot;gelu_pytorch_tanh&quot;)
    self.mlp_residual = ScaleResidual()

    self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
    if self.use_memory:
        self.cam_injector_layer1 = nn.Linear(dim, dim)
        self.cam_injector_layer2 = nn.Linear(dim, dim)
        self.cam_scale_layer = nn.Linear(dim, dim)
        self.cam_shift_layer = nn.Linear(dim, dim)

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    freqs: torch.Tensor,
    # Action Module specific args
    grid_sizes: tuple&#91;int, int, int&#93;,
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    plucker_emb: torch.Tensor | None = None,
    memory_length: int = 0,
    memory_latent_idx: list&#91;int&#93; | tuple&#91;int, ...&#93; | torch.Tensor | None = None,
    predict_latent_idx: tuple&#91;int, int&#93; | list&#91;int&#93; | tuple&#91;int, ...&#93; | torch.Tensor | None = None,
) -&gt; torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)
    orig_dtype = hidden_states.dtype

    if temb.dim() == 4:
        # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
        (
            shift_msa,
            scale_msa,
            gate_msa,
            c_shift_msa,
            c_scale_msa,
            c_gate_msa,
        ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
        # batch_size, seq_len, 1, inner_dim
        shift_msa = shift_msa.squeeze(2)
        scale_msa = scale_msa.squeeze(2)
        gate_msa = gate_msa.squeeze(2)
        c_shift_msa = c_shift_msa.squeeze(2)
        c_scale_msa = c_scale_msa.squeeze(2)
        c_gate_msa = c_gate_msa.squeeze(2)
    else:
        # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
        e = self.scale_shift_table + temb.float()
        (
            shift_msa,
            scale_msa,
            gate_msa,
            c_shift_msa,
            c_scale_msa,
            c_gate_msa,
        ) = e.chunk(6, dim=1)
    assert shift_msa.dtype == torch.float32

    # 1. Self-attention
    norm_hidden_states = (self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa).to(orig_dtype)
    query, _ = self.to_q(norm_hidden_states)
    key, _ = self.to_k(norm_hidden_states)
    value, _ = self.to_v(norm_hidden_states)

    if self.norm_q is not None:
        query = self.norm_q(query)
    if self.norm_k is not None:
        key = self.norm_k(key)

    query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

    # Apply rotary embeddings
    grid_frames, grid_height, grid_width = grid_sizes

    if memory_length &gt; 0:
        hw = grid_height * grid_width
        query_memory = query&#91;:, :memory_length * hw&#93;
        key_memory = key&#91;:, :memory_length * hw&#93;
        query_pred = query&#91;:, memory_length * hw:&#93;
        key_pred = key&#91;:, memory_length * hw:&#93;

        # Indices are pre-built device tensors by MatrixGame3WanModel.forward.
        mem_indices = memory_latent_idx
        if mem_indices is None:
            mem_indices = torch.arange(memory_length, device=query.device, dtype=torch.long)
        pred_indices = predict_latent_idx
        if pred_indices is None:
            pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long)

        query_memory = _apply_rope_with_frame_indices(query_memory,
                                                      freqs,
                                                      height=grid_height,
                                                      width=grid_width,
                                                      frame_indices=mem_indices)
        key_memory = _apply_rope_with_frame_indices(key_memory,
                                                    freqs,
                                                    height=grid_height,
                                                    width=grid_width,
                                                    frame_indices=mem_indices)
        query_pred = _apply_rope_with_frame_indices(query_pred,
                                                    freqs,
                                                    height=grid_height,
                                                    width=grid_width,
                                                    frame_indices=pred_indices)
        key_pred = _apply_rope_with_frame_indices(key_pred,
                                                  freqs,
                                                  height=grid_height,
                                                  width=grid_width,
                                                  frame_indices=pred_indices)
        query = torch.cat(&#91;query_memory, query_pred&#93;, dim=1)
        key = torch.cat(&#91;key_memory, key_pred&#93;, dim=1)
    else:
        pred_indices = predict_latent_idx
        if pred_indices is None:
            pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long)
        query = _apply_rope_with_frame_indices(query,
                                               freqs,
                                               height=grid_height,
                                               width=grid_width,
                                               frame_indices=pred_indices)
        key = _apply_rope_with_frame_indices(key,
                                             freqs,
                                             height=grid_height,
                                             width=grid_width,
                                             frame_indices=pred_indices)

    attn_output, _ = self.attn1(query, key, value)
    attn_output = attn_output.flatten(2)
    attn_output = attn_output.to(self.to_out.weight.dtype)
    attn_output, _ = self.to_out(attn_output)
    attn_output = attn_output.squeeze(1)

    hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float()
    if self.use_memory and plucker_emb is not None:
        c2ws_hidden_states = self.cam_injector_layer2(
            torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)))
        c2ws_hidden_states = c2ws_hidden_states + plucker_emb
        cam_scale = self.cam_scale_layer(c2ws_hidden_states)
        cam_shift = self.cam_shift_layer(c2ws_hidden_states)
        hidden_states = (1.0 + cam_scale) * hidden_states + cam_shift
    norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)

    # 2. Cross-attention
    attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)

    # - use_memory / action path: x = norm3(x); x = x + cross_attn(x)
    # - plain path:              x = x + cross_attn(norm3(x))
    cross_attn_residual_base = (norm_hidden_states if
                                (mouse_cond is not None or self.use_memory) else hidden_states)
    hidden_states = cross_attn_residual_base + attn_output
    norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
    norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa

    # ================= Action Module =================
    if self.action_model is not None and (mouse_cond is not None or keyboard_cond is not None):
        action_dtype = self.ffn.fc_in.weight.dtype
        # grid_sizes is expected to be &#91;F, H, W&#93;
        hidden_states = self.action_model(
            hidden_states.to(action_dtype),
            grid_sizes&#91;0&#93;,
            grid_sizes&#91;1&#93;,
            grid_sizes&#91;2&#93;,
            mouse_cond,
            keyboard_cond,
            mouse_cond_memory=mouse_cond_memory,
            keyboard_cond_memory=keyboard_cond_memory,
        )

        norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
        norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa
    # =================================================

    # 3. Feed-forward
    ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype))
    hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
    hidden_states = hidden_states.to(orig_dtype)

    return hidden_states</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3TransformerBlock.forward
forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs: torch.Tensor, grid_sizes: tuple[int, int, int], mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, plucker_emb: torch.Tensor | None = None, memory_length: int = 0, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    freqs: torch.Tensor,
    # Action Module specific args
    grid_sizes: tuple[int, int, int],
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    plucker_emb: torch.Tensor | None = None,
    memory_length: int = 0,
    memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
    predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
) -> torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)
    orig_dtype = hidden_states.dtype
if temb.dim() == 4:
    # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
    (
        shift_msa,
        scale_msa,
        gate_msa,
        c_shift_msa,
        c_scale_msa,
        c_gate_msa,
    ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
    # batch_size, seq_len, 1, inner_dim
    shift_msa = shift_msa.squeeze(2)
    scale_msa = scale_msa.squeeze(2)
    gate_msa = gate_msa.squeeze(2)
    c_shift_msa = c_shift_msa.squeeze(2)
    c_scale_msa = c_scale_msa.squeeze(2)
    c_gate_msa = c_gate_msa.squeeze(2)
else:
    # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
    e = self.scale_shift_table + temb.float()
    (
        shift_msa,
        scale_msa,
        gate_msa,
        c_shift_msa,
        c_scale_msa,
        c_gate_msa,
    ) = e.chunk(6, dim=1)
assert shift_msa.dtype == torch.float32

# 1. Self-attention
norm_hidden_states = (self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa).to(orig_dtype)
query, _ = self.to_q(norm_hidden_states)
key, _ = self.to_k(norm_hidden_states)
value, _ = self.to_v(norm_hidden_states)

if self.norm_q is not None:
    query = self.norm_q(query)
if self.norm_k is not None:
    key = self.norm_k(key)

query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

# Apply rotary embeddings
grid_frames, grid_height, grid_width = grid_sizes

if memory_length &gt; 0:
    hw = grid_height * grid_width
    query_memory = query&#91;:, :memory_length * hw&#93;
    key_memory = key&#91;:, :memory_length * hw&#93;
    query_pred = query&#91;:, memory_length * hw:&#93;
    key_pred = key&#91;:, memory_length * hw:&#93;

    # Indices are pre-built device tensors by MatrixGame3WanModel.forward.
    mem_indices = memory_latent_idx
    if mem_indices is None:
        mem_indices = torch.arange(memory_length, device=query.device, dtype=torch.long)
    pred_indices = predict_latent_idx
    if pred_indices is None:
        pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long)

    query_memory = _apply_rope_with_frame_indices(query_memory,
                                                  freqs,
                                                  height=grid_height,
                                                  width=grid_width,
                                                  frame_indices=mem_indices)
    key_memory = _apply_rope_with_frame_indices(key_memory,
                                                freqs,
                                                height=grid_height,
                                                width=grid_width,
                                                frame_indices=mem_indices)
    query_pred = _apply_rope_with_frame_indices(query_pred,
                                                freqs,
                                                height=grid_height,
                                                width=grid_width,
                                                frame_indices=pred_indices)
    key_pred = _apply_rope_with_frame_indices(key_pred,
                                              freqs,
                                              height=grid_height,
                                              width=grid_width,
                                              frame_indices=pred_indices)
    query = torch.cat(&#91;query_memory, query_pred&#93;, dim=1)
    key = torch.cat(&#91;key_memory, key_pred&#93;, dim=1)
else:
    pred_indices = predict_latent_idx
    if pred_indices is None:
        pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long)
    query = _apply_rope_with_frame_indices(query,
                                           freqs,
                                           height=grid_height,
                                           width=grid_width,
                                           frame_indices=pred_indices)
    key = _apply_rope_with_frame_indices(key,
                                         freqs,
                                         height=grid_height,
                                         width=grid_width,
                                         frame_indices=pred_indices)

attn_output, _ = self.attn1(query, key, value)
attn_output = attn_output.flatten(2)
attn_output = attn_output.to(self.to_out.weight.dtype)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)

hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float()
if self.use_memory and plucker_emb is not None:
    c2ws_hidden_states = self.cam_injector_layer2(
        torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)))
    c2ws_hidden_states = c2ws_hidden_states + plucker_emb
    cam_scale = self.cam_scale_layer(c2ws_hidden_states)
    cam_shift = self.cam_shift_layer(c2ws_hidden_states)
    hidden_states = (1.0 + cam_scale) * hidden_states + cam_shift
norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)

# 2. Cross-attention
attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)

# - use_memory / action path: x = norm3(x); x = x + cross_attn(x)
# - plain path:              x = x + cross_attn(norm3(x))
cross_attn_residual_base = (norm_hidden_states if
                            (mouse_cond is not None or self.use_memory) else hidden_states)
hidden_states = cross_attn_residual_base + attn_output
norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa

# ================= Action Module =================
if self.action_model is not None and (mouse_cond is not None or keyboard_cond is not None):
    action_dtype = self.ffn.fc_in.weight.dtype
    # grid_sizes is expected to be &#91;F, H, W&#93;
    hidden_states = self.action_model(
        hidden_states.to(action_dtype),
        grid_sizes&#91;0&#93;,
        grid_sizes&#91;1&#93;,
        grid_sizes&#91;2&#93;,
        mouse_cond,
        keyboard_cond,
        mouse_cond_memory=mouse_cond_memory,
        keyboard_cond_memory=keyboard_cond_memory,
    )

    norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
    norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa
# =================================================

# 3. Feed-forward
ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype))
hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
hidden_states = hidden_states.to(orig_dtype)

return hidden_states</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3WanModel
MatrixGame3WanModel(config: MatrixGame3WanVideoConfig, hf_config: dict[str, Any], **kwargs)
Source
class MatrixGame3WanModel(BaseDiT):
    # Marker for action input support (Matrix-Game)
    supports_action_input = True
_fsdp_shard_conditions = _DEFAULT_MATRIXGAME3_CONFIG._fsdp_shard_conditions
_compile_conditions = _DEFAULT_MATRIXGAME3_CONFIG._compile_conditions
_supported_attention_backends = (_DEFAULT_MATRIXGAME3_CONFIG._supported_attention_backends)
param_names_mapping = _DEFAULT_MATRIXGAME3_CONFIG.param_names_mapping
reverse_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.reverse_param_names_mapping)
lora_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.lora_param_names_mapping)

def __init__(
    self,
    config: MatrixGame3WanVideoConfig,
    hf_config: dict&#91;str, Any&#93;,
    **kwargs,
) -&gt; None:
    super().__init__(config=config, hf_config=hf_config)

    inner_dim = config.num_attention_heads * config.attention_head_dim
    self.hidden_size = config.hidden_size
    self.num_attention_heads = config.num_attention_heads
    self.num_channels_latents = config.out_channels
    self.patch_size = config.patch_size

    # 1. Patch &amp; position embedding
    self.patch_embedding = PatchEmbed(
        in_chans=config.in_channels,
        embed_dim=inner_dim,
        patch_size=config.patch_size,
        flatten=False,
    )

    # 2. Condition embeddings
    self.condition_embedder = MatrixGame3TimeImageEmbedding(
        dim=inner_dim,
        time_freq_dim=config.freq_dim,
    )
    self.text_embedding = nn.Sequential(
        nn.Linear(config.text_dim, inner_dim),
        nn.GELU(approximate=&quot;tanh&quot;),
        nn.Linear(inner_dim, inner_dim),
    )
    self.use_memory = getattr(config, &quot;use_memory&quot;, False)
    self.sigma_theta = float(getattr(config, &quot;sigma_theta&quot;, 0.0))
    self.camera_patch_embedding = None
    self.c2ws_hidden_states_layer1 = None
    self.c2ws_hidden_states_layer2 = None
    if self.use_memory:
        camera_in_channels = getattr(config, &quot;camera_embed_in_channels&quot;, 1536)
        self.camera_patch_embedding = WanCamControlPatchEmbedding(
            patch_size=config.patch_size,
            in_chans=camera_in_channels,
            embed_dim=inner_dim,
        )
        self.c2ws_hidden_states_layer1 = nn.Linear(inner_dim, inner_dim)
        self.c2ws_hidden_states_layer2 = nn.Linear(inner_dim, inner_dim)

    # 2.1. Get action config
    self.action_config = getattr(config, &quot;action_config&quot;, &#123;&#125;)

    # 3. Transformer blocks
    self.blocks = nn.ModuleList(&#91;
        MatrixGame3TransformerBlock(
            inner_dim,
            config.ffn_dim,
            config.num_attention_heads,
            config.qk_norm,
            config.cross_attn_norm,
            config.eps,
            self._supported_attention_backends,
            prefix=f&quot;&#123;getattr(config, &#x27;prefix&#x27;, &#x27;Wan&#x27;)&#125;.blocks.&#123;i&#125;&quot;,
            action_config=self.action_config,
            block_id=i,
            use_memory=self.use_memory,
        ) for i in range(config.num_layers)
    &#93;)

    # 4. Output norm &amp; projection
    self.norm_out = LayerNormScaleShift(
        inner_dim,
        norm_type=&quot;layer&quot;,
        eps=config.eps,
        elementwise_affine=False,
        dtype=torch.float32,
        compute_dtype=torch.float32,
    )
    self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size))
    self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)

    self.gradient_checkpointing = False

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list&#91;torch.Tensor&#93;,
    timestep: torch.LongTensor,
    # Action inputs
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    x_memory: torch.Tensor | None = None,
    timestep_memory: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    c2ws_plucker_emb: torch.Tensor | None = None,
    memory_latent_idx: list&#91;int&#93; | tuple&#91;int, ...&#93; | torch.Tensor | None = None,
    predict_latent_idx: tuple&#91;int, int&#93; | list&#91;int&#93; | tuple&#91;int, ...&#93; | torch.Tensor | None = None,
    **kwargs,
) -&gt; torch.Tensor:
    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;

    memory_length = 0
    if x_memory is not None:
        memory_length = x_memory.shape&#91;2&#93;
        hidden_states = torch.cat(&#91;x_memory.to(hidden_states.dtype), hidden_states&#93;, dim=2)

    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w
    max_frame_index = post_patch_num_frames
    if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
        max_frame_index = max(max_frame_index, int(predict_latent_idx&#91;1&#93;))
    elif predict_latent_idx is not None:
        max_frame_index = max(max_frame_index, int(max(predict_latent_idx)) + 1)
    if memory_latent_idx is not None and len(memory_latent_idx) &gt; 0:
        max_frame_index = max(max_frame_index, int(max(memory_latent_idx)) + 1)

    # Pre-build RoPE frame index tensors once per forward so blocks don&#x27;t
    # do `torch.tensor(list, device=cuda)` per call (host sync).
    device_idx = hidden_states.device
    if predict_latent_idx is None:
        predict_latent_idx = None
    elif torch.is_tensor(predict_latent_idx):
        predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long)
    elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
        predict_latent_idx = torch.arange(
            predict_latent_idx&#91;0&#93;,
            predict_latent_idx&#91;1&#93;,
            device=device_idx,
            dtype=torch.long,
        )
    else:
        predict_latent_idx = torch.tensor(
            list(predict_latent_idx),
            device=device_idx,
            dtype=torch.long,
        )

    if memory_latent_idx is None or (not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0):
        memory_latent_idx = None
    elif torch.is_tensor(memory_latent_idx):
        memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long)
    else:
        memory_latent_idx = torch.tensor(
            list(memory_latent_idx),
            device=device_idx,
            dtype=torch.long,
        )

    head_dim = self.hidden_size // self.num_attention_heads
    freqs = _build_rope_freqs(
        max_seq_len=(2048 if self.use_memory else 1024),
        head_dim=head_dim,
        num_heads=self.num_attention_heads,
        sigma_theta=(self.sigma_theta if self.use_memory else 0.0),
        device=hidden_states.device,
    )
    max_seq_len = freqs.shape&#91;1&#93; if freqs.dim() == 3 else freqs.shape&#91;0&#93;
    if max_frame_index &gt; max_seq_len or post_patch_height &gt; max_seq_len or post_patch_width &gt; max_seq_len:
        raise ValueError(f&quot;got frames=&#123;max_frame_index&#125;, height=&#123;post_patch_height&#125;, width=&#123;post_patch_width&#125;, &quot;
                         f&quot;max_seq_len=&#123;max_seq_len&#125;&quot;)

    hidden_states = self.patch_embedding(hidden_states)
    hidden_states = hidden_states.flatten(2).transpose(1, 2)

    plucker_emb = None
    if c2ws_plucker_emb is not None and self.camera_patch_embedding is not None:
        if memory_length &gt; 0 and c2ws_plucker_emb.shape&#91;2&#93; == post_patch_num_frames - memory_length:
            zeros = torch.zeros(c2ws_plucker_emb.shape&#91;0&#93;,
                                c2ws_plucker_emb.shape&#91;1&#93;,
                                memory_length,
                                c2ws_plucker_emb.shape&#91;3&#93;,
                                c2ws_plucker_emb.shape&#91;4&#93;,
                                device=c2ws_plucker_emb.device,
                                dtype=c2ws_plucker_emb.dtype)
            c2ws_plucker_emb = torch.cat(&#91;zeros, c2ws_plucker_emb&#93;, dim=2)
        target_frames = post_patch_num_frames * p_t
        target_height = post_patch_height * p_h
        target_width = post_patch_width * p_w
        if c2ws_plucker_emb.shape&#91;2&#93; &gt; target_frames or c2ws_plucker_emb.shape&#91;
                3&#93; &gt; target_height or c2ws_plucker_emb.shape&#91;4&#93; &gt; target_width:
            c2ws_plucker_emb = c2ws_plucker_emb&#91;:, :, :target_frames, :target_height, :target_width&#93;
        if c2ws_plucker_emb.shape&#91;2:&#93; != (target_frames, target_height, target_width):
            raise ValueError(&quot;Camera conditioning must align with the patchified latent grid, &quot;
                             f&quot;got camera shape &#123;tuple(c2ws_plucker_emb.shape&#91;2:&#93;)&#125; vs target &quot;
                             f&quot;&#123;(target_frames, target_height, target_width)&#125;&quot;)
        plucker_emb = self.camera_patch_embedding(c2ws_plucker_emb.to(hidden_states.dtype))
        plucker_hidden = self.c2ws_hidden_states_layer2(
            torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)))
        plucker_emb = plucker_emb + plucker_hidden

    timestep_tokens = timestep
    if timestep_tokens.dim() == 0:
        timestep_tokens = timestep_tokens.unsqueeze(0)
    if timestep_tokens.dim() == 1:
        timestep_tokens = timestep_tokens.unsqueeze(1).repeat(
            1, post_patch_num_frames * post_patch_height * post_patch_width)
    elif timestep_tokens.dim() == 2 and timestep_tokens.shape&#91;1&#93; == num_frames:
        timestep_tokens = (timestep_tokens&#91;:, :, None,
                                           None&#93;.repeat(1, 1, post_patch_height,
                                                        post_patch_width).reshape(timestep_tokens.shape&#91;0&#93;, -1))
    if memory_length &gt; 0:
        if timestep_memory is None:
            raise ValueError(&quot;MatrixGame3 requires timestep_memory when x_memory is provided&quot;)
        timestep_tokens = torch.cat(&#91;timestep_memory.to(timestep_tokens.dtype), timestep_tokens&#93;, dim=1)

    temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
        timestep_tokens.flatten(),
        encoder_hidden_states,
        timestep_seq_len=timestep_tokens.shape&#91;1&#93;,
    )
    if timestep_proj.dim() == 3 and timestep_proj.shape&#91;-1&#93; % 6 == 0:
        timestep_proj = timestep_proj.unflatten(2, (6, -1))

    if encoder_hidden_states is not None:
        if isinstance(encoder_hidden_states, list):
            encoder_hidden_states = encoder_hidden_states&#91;0&#93;
        elif encoder_hidden_states.ndim == 2:
            encoder_hidden_states = encoder_hidden_states.unsqueeze(0)
        if (encoder_hidden_states.ndim == 3
                and encoder_hidden_states.shape&#91;-1&#93; == self.text_embedding&#91;0&#93;.in_features):
            encoder_hidden_states = self.text_embedding(encoder_hidden_states.to(hidden_states.dtype))

    # &#91;F, H, W&#93; for the ActionModule
    grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)

    # Blocks
    for block in self.blocks:
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            hidden_states = self._gradient_checkpointing_func(
                block,
                hidden_states,
                encoder_hidden_states,
                timestep_proj,
                freqs,
                grid_sizes=grid_sizes,
                mouse_cond=mouse_cond,
                keyboard_cond=keyboard_cond,
                mouse_cond_memory=mouse_cond_memory,
                keyboard_cond_memory=keyboard_cond_memory,
                plucker_emb=plucker_emb,
                memory_length=memory_length,
                memory_latent_idx=memory_latent_idx,
                predict_latent_idx=predict_latent_idx,
            )
        else:
            hidden_states = block(
                hidden_states,
                encoder_hidden_states,
                timestep_proj,
                freqs,
                grid_sizes=grid_sizes,
                mouse_cond=mouse_cond,
                keyboard_cond=keyboard_cond,
                mouse_cond_memory=mouse_cond_memory,
                keyboard_cond_memory=keyboard_cond_memory,
                plucker_emb=plucker_emb,
                memory_length=memory_length,
                memory_latent_idx=memory_latent_idx,
                predict_latent_idx=predict_latent_idx,
            )

    # Output
    if temb.dim() == 3:
        shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
        shift = shift.squeeze(2)
        scale = scale.squeeze(2)
    else:
        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
    hidden_states = self.norm_out(hidden_states, shift, scale)
    hidden_states = self.proj_out(hidden_states)

    hidden_states = hidden_states.reshape(
        batch_size,
        post_patch_num_frames,
        post_patch_height,
        post_patch_width,
        p_t,
        p_h,
        p_w,
        -1,
    )
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

    if memory_length &gt; 0:
        output = output&#91;:, :, memory_length:&#93;

    return output</code></pre></details>
fastvideo.models.dits.matrixgame3.model.MatrixGame3WanModel.forward
forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, x_memory: torch.Tensor | None = None, timestep_memory: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, c2ws_plucker_emb: torch.Tensor | None = None, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, **kwargs) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list[torch.Tensor],
    timestep: torch.LongTensor,
    # Action inputs
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    x_memory: torch.Tensor | None = None,
    timestep_memory: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    c2ws_plucker_emb: torch.Tensor | None = None,
    memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
    predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
    **kwargs,
) -> torch.Tensor:
    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]
memory_length = 0
if x_memory is not None:
    memory_length = x_memory.shape&#91;2&#93;
    hidden_states = torch.cat(&#91;x_memory.to(hidden_states.dtype), hidden_states&#93;, dim=2)

batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.patch_size
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p_h
post_patch_width = width // p_w
max_frame_index = post_patch_num_frames
if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
    max_frame_index = max(max_frame_index, int(predict_latent_idx&#91;1&#93;))
elif predict_latent_idx is not None:
    max_frame_index = max(max_frame_index, int(max(predict_latent_idx)) + 1)
if memory_latent_idx is not None and len(memory_latent_idx) &gt; 0:
    max_frame_index = max(max_frame_index, int(max(memory_latent_idx)) + 1)

# Pre-build RoPE frame index tensors once per forward so blocks don&#x27;t
# do `torch.tensor(list, device=cuda)` per call (host sync).
device_idx = hidden_states.device
if predict_latent_idx is None:
    predict_latent_idx = None
elif torch.is_tensor(predict_latent_idx):
    predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long)
elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
    predict_latent_idx = torch.arange(
        predict_latent_idx&#91;0&#93;,
        predict_latent_idx&#91;1&#93;,
        device=device_idx,
        dtype=torch.long,
    )
else:
    predict_latent_idx = torch.tensor(
        list(predict_latent_idx),
        device=device_idx,
        dtype=torch.long,
    )

if memory_latent_idx is None or (not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0):
    memory_latent_idx = None
elif torch.is_tensor(memory_latent_idx):
    memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long)
else:
    memory_latent_idx = torch.tensor(
        list(memory_latent_idx),
        device=device_idx,
        dtype=torch.long,
    )

head_dim = self.hidden_size // self.num_attention_heads
freqs = _build_rope_freqs(
    max_seq_len=(2048 if self.use_memory else 1024),
    head_dim=head_dim,
    num_heads=self.num_attention_heads,
    sigma_theta=(self.sigma_theta if self.use_memory else 0.0),
    device=hidden_states.device,
)
max_seq_len = freqs.shape&#91;1&#93; if freqs.dim() == 3 else freqs.shape&#91;0&#93;
if max_frame_index &gt; max_seq_len or post_patch_height &gt; max_seq_len or post_patch_width &gt; max_seq_len:
    raise ValueError(f&quot;got frames=&#123;max_frame_index&#125;, height=&#123;post_patch_height&#125;, width=&#123;post_patch_width&#125;, &quot;
                     f&quot;max_seq_len=&#123;max_seq_len&#125;&quot;)

hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)

plucker_emb = None
if c2ws_plucker_emb is not None and self.camera_patch_embedding is not None:
    if memory_length &gt; 0 and c2ws_plucker_emb.shape&#91;2&#93; == post_patch_num_frames - memory_length:
        zeros = torch.zeros(c2ws_plucker_emb.shape&#91;0&#93;,
                            c2ws_plucker_emb.shape&#91;1&#93;,
                            memory_length,
                            c2ws_plucker_emb.shape&#91;3&#93;,
                            c2ws_plucker_emb.shape&#91;4&#93;,
                            device=c2ws_plucker_emb.device,
                            dtype=c2ws_plucker_emb.dtype)
        c2ws_plucker_emb = torch.cat(&#91;zeros, c2ws_plucker_emb&#93;, dim=2)
    target_frames = post_patch_num_frames * p_t
    target_height = post_patch_height * p_h
    target_width = post_patch_width * p_w
    if c2ws_plucker_emb.shape&#91;2&#93; &gt; target_frames or c2ws_plucker_emb.shape&#91;
            3&#93; &gt; target_height or c2ws_plucker_emb.shape&#91;4&#93; &gt; target_width:
        c2ws_plucker_emb = c2ws_plucker_emb&#91;:, :, :target_frames, :target_height, :target_width&#93;
    if c2ws_plucker_emb.shape&#91;2:&#93; != (target_frames, target_height, target_width):
        raise ValueError(&quot;Camera conditioning must align with the patchified latent grid, &quot;
                         f&quot;got camera shape &#123;tuple(c2ws_plucker_emb.shape&#91;2:&#93;)&#125; vs target &quot;
                         f&quot;&#123;(target_frames, target_height, target_width)&#125;&quot;)
    plucker_emb = self.camera_patch_embedding(c2ws_plucker_emb.to(hidden_states.dtype))
    plucker_hidden = self.c2ws_hidden_states_layer2(
        torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)))
    plucker_emb = plucker_emb + plucker_hidden

timestep_tokens = timestep
if timestep_tokens.dim() == 0:
    timestep_tokens = timestep_tokens.unsqueeze(0)
if timestep_tokens.dim() == 1:
    timestep_tokens = timestep_tokens.unsqueeze(1).repeat(
        1, post_patch_num_frames * post_patch_height * post_patch_width)
elif timestep_tokens.dim() == 2 and timestep_tokens.shape&#91;1&#93; == num_frames:
    timestep_tokens = (timestep_tokens&#91;:, :, None,
                                       None&#93;.repeat(1, 1, post_patch_height,
                                                    post_patch_width).reshape(timestep_tokens.shape&#91;0&#93;, -1))
if memory_length &gt; 0:
    if timestep_memory is None:
        raise ValueError(&quot;MatrixGame3 requires timestep_memory when x_memory is provided&quot;)
    timestep_tokens = torch.cat(&#91;timestep_memory.to(timestep_tokens.dtype), timestep_tokens&#93;, dim=1)

temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
    timestep_tokens.flatten(),
    encoder_hidden_states,
    timestep_seq_len=timestep_tokens.shape&#91;1&#93;,
)
if timestep_proj.dim() == 3 and timestep_proj.shape&#91;-1&#93; % 6 == 0:
    timestep_proj = timestep_proj.unflatten(2, (6, -1))

if encoder_hidden_states is not None:
    if isinstance(encoder_hidden_states, list):
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;
    elif encoder_hidden_states.ndim == 2:
        encoder_hidden_states = encoder_hidden_states.unsqueeze(0)
    if (encoder_hidden_states.ndim == 3
            and encoder_hidden_states.shape&#91;-1&#93; == self.text_embedding&#91;0&#93;.in_features):
        encoder_hidden_states = self.text_embedding(encoder_hidden_states.to(hidden_states.dtype))

# &#91;F, H, W&#93; for the ActionModule
grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)

# Blocks
for block in self.blocks:
    if torch.is_grad_enabled() and self.gradient_checkpointing:
        hidden_states = self._gradient_checkpointing_func(
            block,
            hidden_states,
            encoder_hidden_states,
            timestep_proj,
            freqs,
            grid_sizes=grid_sizes,
            mouse_cond=mouse_cond,
            keyboard_cond=keyboard_cond,
            mouse_cond_memory=mouse_cond_memory,
            keyboard_cond_memory=keyboard_cond_memory,
            plucker_emb=plucker_emb,
            memory_length=memory_length,
            memory_latent_idx=memory_latent_idx,
            predict_latent_idx=predict_latent_idx,
        )
    else:
        hidden_states = block(
            hidden_states,
            encoder_hidden_states,
            timestep_proj,
            freqs,
            grid_sizes=grid_sizes,
            mouse_cond=mouse_cond,
            keyboard_cond=keyboard_cond,
            mouse_cond_memory=mouse_cond_memory,
            keyboard_cond_memory=keyboard_cond_memory,
            plucker_emb=plucker_emb,
            memory_length=memory_length,
            memory_latent_idx=memory_latent_idx,
            predict_latent_idx=predict_latent_idx,
        )

# Output
if temb.dim() == 3:
    shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
    shift = shift.squeeze(2)
    scale = scale.squeeze(2)
else:
    shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states, shift, scale)
hidden_states = self.proj_out(hidden_states)

hidden_states = hidden_states.reshape(
    batch_size,
    post_patch_num_frames,
    post_patch_height,
    post_patch_width,
    p_t,
    p_h,
    p_w,
    -1,
)
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

if memory_length &gt; 0:
    output = output&#91;:, :, memory_length:&#93;

return output</code></pre></details>