def apply_audio_delay(
    audio_BxTxC: mx.array,
    pad_value: int,
    bos_value: int,
    precomp: tp.Tuple[mx.array, mx.array],
) -> mx.array:
    """
    Applies the delay pattern to batched audio tokens using precomputed indices,
    inserting BOS where t_idx < 0 and PAD where t_idx >= T.

    Args:
        audio_BxTxC: [B, T, C] int16 audio tokens (or int32/float)
        pad_value: the padding token
        bos_value: the BOS token
        precomp:  (t_idx_BxTxC, indices_BTCx3) from build_delay_indices

    Returns:
        result_BxTxC: [B, T, C] delayed audio tokens
    """
    t_idx_BxTxC, indices_BTCx3 = precomp

    def gather_nd(array, indices):
        gathered = []
        for idx in range(indices.shape[0]):
            b, t, c = indices[idx, 0], indices[idx, 1], indices[idx, 2]
            gathered.append(array[b, t, c])
        return mx.array(gathered)

    # Apply gather
    gathered_flat = gather_nd(audio_BxTxC, indices_BTCx3)
    gathered_BxTxC = mx.reshape(gathered_flat, audio_BxTxC.shape)

    # Create masks
    mask_bos = t_idx_BxTxC < 0  # => place bos_value
    mask_pad = t_idx_BxTxC >= audio_BxTxC.shape[1]  # => place pad_value

    # Create scalar values
    bos_tensor = mx.full(1, bos_value, dtype=audio_BxTxC.dtype)
    pad_tensor = mx.full(1, pad_value, dtype=audio_BxTxC.dtype)

    # Apply masks (if mask_bos, BOS; else if mask_pad, PAD; else original gather)
    result_BxTxC = mx.where(
        mask_bos, bos_tensor, mx.where(mask_pad, pad_tensor, gathered_BxTxC)
    )

    return result_BxTxC

    Applies the delay pattern to batched audio tokens using precomputed indices,
    inserting BOS where t_idx < 0 and PAD where t_idx >= T.

    Args:
        audio_BxTxC: [B, T, C] int16 audio tokens (or int32/float)
        pad_value: the padding token
        bos_value: the BOS token
        precomp:  (t_idx_BxTxC, indices_BTCx3) from build_delay_indices

    Returns:
        result_BxTxC: [B, T, C] delayed audio tokens
    
mlx_audio/tts/models/dia/audio.py