def PhiOAudioEmbedding_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, torch.nn.Sequential):
        if self.audio_projection[0].weight.dtype in [
            torch.qint8,
            torch.int8,
            torch.uint8,
        ]:
            target_dtype = self.audio_projection[0]._op_context.get_bias().dtype
        else:
            target_dtype = self.audio_projection[0].bias.dtype
    elif isinstance(self.audio_projection, torch.nn.ModuleDict):
        if self.audio_projection[audio_projection_mode][0].weight.dtype in [
            torch.qint8,
            torch.int8,
            torch.uint8,
        ]:
            target_dtype = (
                self.audio_projection[audio_projection_mode][0]
                ._op_context.get_bias()
                .dtype
            )
        else:
            target_dtype = self.audio_projection[audio_projection_mode][0].bias.dtype
    else:  # It's a single nn.Linear layer
        if self.audio_projection.weight.dtype in [torch.qint8, torch.int8, torch.uint8]:
            target_dtype = self.audio_projection._op_context.get_bias().dtype
        else:
            target_dtype = self.audio_projection.bias.dtype

    if input_embeds is not None:
        input_embeds = input_embeds.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_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
        ), f"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}"

        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
    
intel_extension_for_pytorch/transformers/models/reference/models.py