Question: Can modality row selection be moved before the output projections?

Question

In MiniMaxH3Transformer3DModel.forward, the output heads are currently applied to the entire packed sequence before selecting the video and audio rows:

hidden_states = self.norm_out(hidden_states, temb, timestep_indices)

video_output = self.proj_out(hidden_states).index_select(1, video_indices)
audio_output = self.audio_proj_out(hidden_states).index_select(1, audio_indices)

Since both output projections operate independently on the last feature dimension, would the following be mathematically equivalent?

hidden_states = self.norm_out(hidden_states, temb, timestep_indices)

video_hidden_states = hidden_states.index_select(1, video_indices)
audio_hidden_states = hidden_states.index_select(1, audio_indices)

video_output = self.proj_out(video_hidden_states)
audio_output = self.audio_proj_out(audio_hidden_states)

In other words:
index_select(Linear(H), I) = Linear(index_select(H, I))
This would avoid applying proj_out and audio_proj_out to rows that are not needed by the corresponding output head.
Is there any reason why the current order is required, such as:

• context-parallel output gathering;
• global versus local sequence indices;
• distributed inference constraints;
• numerical or checkpoint compatibility;
• or the need to preserve predictions for all reference rows?
  I would appreciate any clarification about whether selecting the rows before the modality-specific output projections is a valid optimization.

Start at MiniMaxH3Transformer3DModel.forward and trace how norm_out, proj_out, audio_proj_out, video_indices, and audio_indices interact with context-parallel output gathering. Check whether indices are local or global and whether distributed inference or reference-row predictions depend on the current order; done means establishing the optimization's validity and documenting any constraints.