def forward(
        self,
        input_ids: torch.LongTensor,
        input_embeds: torch.FloatTensor,
        audio_embed_sizes=None,
        audio_attention_mask=None,
        audio_projection_mode="speech",
        **kwargs,
    ) -> torch.FloatTensor:
        """
        arguments:
            input_ids: input text ids (B, U)
            input_embeds: audio features (B, T, D)  B: num audios in a sequence
        """
        if self.input_embeds is not None:
            input_embeds = self.input_embeds.clone()
        if self.audio_embed_sizes is not None:
            audio_embed_sizes = self.audio_embed_sizes.clone()

        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])
        MAX_INPUT_ID = int(1e9)

        with torch.no_grad():
            positions = torch.nonzero(
                input_ids == _AUDIO_SPECIAL_TOKEN_ID, as_tuple=False
            )
            positions_tuple = torch.nonzero(
                input_ids == _AUDIO_SPECIAL_TOKEN_ID, as_tuple=True
            )

        if isinstance(self.audio_projection, nn.Sequential):
            target_device = self.audio_projection[0].bias.device
            target_dtype = self.audio_projection[0].bias.dtype
        elif isinstance(self.audio_projection, nn.ModuleDict):
            target_device = self.audio_projection[audio_projection_mode][0].bias.device
            target_dtype = self.audio_projection[audio_projection_mode][0].bias.dtype
        else:  # It's a single nn.Linear layer
            target_device = self.audio_projection.bias.device
            target_dtype = self.audio_projection.bias.dtype

        if input_embeds is not None:
            input_embeds = input_embeds.to(target_device).to(target_dtype)

        if len(positions.tolist()) > 0:
            audio_set_tensor = self.get_audio_features(
                input_embeds, audio_attention_mask, audio_projection_mode
            )
        else:
            # # create an audio tensor
            # To do: not sure if this is required for text only input
            if self.training:
                audio_embeds = (
                    torch.zeros(1, 500, self.audio_dim_in)
                    .to(target_device)
                    .to(target_dtype)
                )
                audio_attention_mask = audio_embeds.new_ones(
                    audio_embeds.size()[:2]
                ).long()
                audio_set_tensor = self.get_audio_features(
                    audio_embeds, audio_attention_mask, audio_projection_mode
                )

        hidden_states = kwargs["wte"](input_ids)

        if len(positions.tolist()) > 0:

            assert audio_embed_sizes.sum().item() == len(
                positions
            ), "please ensure the encoder outputs have the same length as defined in input_ids! \n \
            audio_embed_sizes.sum().item(): {audio_embed_sizes.sum().item()} \n \
            len(positions): {len(positions)} \n audio_embed_sizes: {audio_embed_sizes} \n \
            positions: {positions} \n input_ids.shape \n {input_ids.shape}"

            # new implementation without in-place operation
            # Ref: https://pytorch.org/docs/stable/generated/torch.Tensor.index_put.html
            # Ref: https://pytorch.org/docs/stable/generated/torch.Tensor.index_put_.html#torch.Tensor.index_put_
            # audio_set_tensor: shape (N_audios, N_padded_tokens, C)
            # Shape: (merged_N_tokens, C)
            merged_audio_set_tensor = torch.cat(
                [
                    audio_set_tensor[i, : audio_embed_sizes[i], :]
                    for i in range(len(audio_embed_sizes))
                ],
                dim=0,
            )
            merged_audio_set_tensor = merged_audio_set_tensor.to(
                hidden_states.dtype
            ).to(hidden_states.device)
            # Temporarily disable autocast to avoid issue on bf16 tensors
            # Ref: https://github.com/pytorch/pytorch/issues/132715
            with torch.autocast(device_type=hidden_states.device.type, enabled=False):
                new_hidden_states = hidden_states.index_put(
                    indices=positions_tuple,
                    values=merged_audio_set_tensor,
                    accumulate=False,
                )
            hidden_states = new_hidden_states
        else:
            if self.training:
                hidden_states = (
                    hidden_states
                    + (
                        0
                        * audio_set_tensor[:, 0]
                        .to(hidden_states.dtype)
                        .to(hidden_states.device)
                    ).sum()
                )

        if self.drop is not None:
            hidden_states = self.drop(hidden_states)

        return hidden_states

        arguments:
            input_ids: input text ids (B, U)
            input_embeds: audio features (B, T, D)  B: num audios in a sequence
        
tests/cpu/hf_configs/phi4/modeling_phi4mm.py