xref: /aosp_15_r20/external/executorch/extension/llm/modules/__init__.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7from ._position_embeddings import (
8    replace_tile_positional_embedding,
9    replace_tiled_token_positional_embedding,
10    TiledTokenPositionalEmbedding,
11    TilePositionalEmbedding,
12)
13from .attention import MultiHeadAttention, replace_mha_with_inference_mha
14from .kv_cache import KVCache
15
16__all__ = [
17    "TilePositionalEmbedding",
18    "TiledTokenPositionalEmbedding",
19    "replace_tile_positional_embedding",
20    "replace_tiled_token_positional_embedding",
21    "MultiHeadAttention",
22    "replace_mha_with_inference_mha",
23    "KVCache",
24]
25