ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -0,0 +1,246 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Attention modules for the T5 encoder-decoder model family.
+
+from typing import Optional, Tuple
+
+from ai_edge_torch.generative.layers.attention import CrossAttention
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.kv_cache import KVCache
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
+import torch
+from torch import nn
+
+BATCH_SIZE = 1
+
+
+class EncoderDecoderBlock(nn.Module):
+
+  def __init__(
+      self,
+      config: cfg.TransformerBlockConfig,
+      model_config: cfg.ModelConfig,
+      has_relative_attention_bias: bool = False,
+  ) -> None:
+    """Initialize an instance of the EncoderDecoderBlock.
+
+    Args:
+      config (cfg.TransformerBlockConfig): the configuration object for this
+        transformer block.
+      model_config (cfg.ModelConfig): the configuration object for the model
+        this transformer block belongs to.
+      has_relative_attention_bias (bool): whether the self attention block has
+        relative bias.
+    """
+
+    super().__init__()
+    self.atten_func = T5Attention(
+        BATCH_SIZE,
+        model_config.embedding_dim,
+        config.attn_config,
+        config.pre_attention_norm_config,
+        model_config.kv_cache_max,
+        model_config.enable_hlfb,
+        has_relative_attention_bias=has_relative_attention_bias,
+    )
+    # For a decoder, we add a cross attention.
+    if model_config.is_decoder:
+      self.cross_atten_func = T5Attention(
+          BATCH_SIZE,
+          model_config.embedding_dim,
+          config.attn_config,
+          config.pre_attention_norm_config,
+          model_config.kv_cache_max,
+          model_config.enable_hlfb,
+          # Cross Attention does not have relative attention bias.
+          has_relative_attention_bias=False,
+      )
+    else:
+      self.cross_atten_func = None
+
+    self.post_atten_norm = builder.build_norm(
+        model_config.embedding_dim,
+        config.post_attention_norm_config,
+    )
+    self.ff = builder.build_ff(model_config.embedding_dim, config.ff_config)
+    self.config = config
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+      encoder_hidden_states: Optional[torch.Tensor] = None,
+      encoder_attention_mask: Optional[torch.Tensor] = None,
+      encoder_decoder_position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the EncoderDecoderBlock.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this transformer block.
+    """
+
+    hidden_states, position_bias = self.atten_func(
+        x,
+        input_pos=input_pos,
+        mask=mask,
+        relative_position=relative_position,
+        position_bias=position_bias,
+    )
+
+    attn_out = hidden_states + x
+
+    if self.cross_atten_func:
+      hidden_states, encoder_decoder_position_bias = self.cross_atten_func(
+          attn_out,
+          input_pos=input_pos,
+          key_value_states=encoder_hidden_states,
+          mask=encoder_attention_mask,
+          relative_position=relative_position,
+          position_bias=encoder_decoder_position_bias,
+      )
+      attn_out = hidden_states + attn_out
+
+    forwarded = self.post_atten_norm(attn_out)
+    forwarded = self.ff(forwarded)
+    hidden_states = attn_out + forwarded
+
+    # encoder_deocder_position_bias is from CrossAttention
+    return hidden_states, position_bias, encoder_decoder_position_bias
+
+
+class T5Attention(CrossAttention):
+
+  def __init__(
+      self,
+      batch: int,
+      dim: int,
+      config: cfg.AttentionConfig,
+      norm_config: cfg.NormalizationConfig,
+      kv_cache_max: int,
+      enable_hlfb: bool,
+      has_relative_attention_bias=False,
+  ) -> None:
+    """Initialize an instance of T5Attention.
+
+    Args:
+      dim (int): causal attention's input/output dimmension.
+      config (cfg.AttentionConfig): attention specific configurations.
+      norm_config (cfg.NormalizationConfig): normalization configure before
+        attention.
+      kv_cache_max (int): determines the size of the KV Cache buffer, if
+        enabled.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+      has_relative_attention_bias (bool): whether we compute relative bias.
+    """
+    super().__init__(batch, dim, dim, config, kv_cache_max, enable_hlfb)
+    self.pre_atten_norm = builder.build_norm(dim, norm_config)
+
+    self.has_relative_attention_bias = has_relative_attention_bias
+    self.relative_attention_num_buckets = config.relative_attention_num_buckets
+    if self.has_relative_attention_bias:
+      self.relative_attention_bias = nn.Embedding(
+          self.relative_attention_num_buckets, self.n_heads
+      )
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      key_value_states: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the T5Attention layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this self attention layer.
+    """
+
+    x = self.pre_atten_norm(x)
+    B, T, C = (
+        x.size()
+    )  # batch size, sequence length, embedding dimensionality (n_embd)
+    query_states = self.q_projection(x)
+    query_states = query_states.reshape(
+        B, T, -1, self.config.head_dim
+    )  # (B, T, nh_q, hs)
+
+    if key_value_states is not None:
+      (
+          kvB,
+          kvT,
+          kvC,
+      ) = (
+          key_value_states.size()
+      )  # batch size, sequence length, embedding dimensionality (n_embd)
+      key_states = self.k_projection(key_value_states)
+      value_states = self.v_projection(key_value_states)
+      key_states = key_states.reshape(kvB, kvT, -1, self.config.head_dim)
+      value_states = value_states.reshape(kvB, kvT, -1, self.config.head_dim)
+    else:
+      key_states = self.k_projection(x)
+      value_states = self.v_projection(x)
+      key_states = key_states.reshape(B, T, -1, self.config.head_dim)
+      value_states = value_states.reshape(B, T, -1, self.config.head_dim)
+
+    if key_value_states is None and self.kv_cache is not None:
+      key_states, value_states = self.kv_cache.update_cache(
+          input_pos, key_states, value_states
+      )
+
+    if position_bias is None:
+      # handle the encoder case first
+      if self.has_relative_attention_bias:
+        position_bias = self.relative_attention_bias(
+            relative_position
+        )  # shape (query_length, key_length, num_heads)
+        position_bias = position_bias.permute([0, 1, 4, 2, 3]).squeeze(
+            0
+        )  # shape (1, num_heads, query_length, key_length)
+      else:
+        # position_bias = torch.zeros(B, self.n_heads, T, self.config.head_dim, dtype=torch.float32)
+        position_bias = torch.zeros_like(mask, dtype=torch.float32)
+
+    mask = mask + position_bias
+    y = self.sdpa_func(
+        query_states,
+        key_states,
+        value_states,
+        self.config.head_dim,
+        mask=mask,
+        scale=1.0,
+    )
+    y = y.reshape(B, T, C)  # re-assemble all head outputs side by side
+    # output projection
+    y = self.output_projection(y)
+    return y, position_bias

ai_edge_torch/generative/examples/test_models/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/generative/examples/test_models/convert_toy_model.py

@@ -0,0 +1,105 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has a single-layer transformer block.
+from absl import app
+import ai_edge_torch
+from ai_edge_torch import lowertools
+from ai_edge_torch.generative.examples.test_models import toy_model
+from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import torch
+
+KV_CACHE_MAX_LEN = 100
+
+
+def convert_toy_model(_) -> None:
+  """Converts a toy model to tflite."""
+  model = toy_model.ToySingleLayerModel(toy_model.get_model_config())
+  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 0)
+  input_pos = torch.arange(0, KV_CACHE_MAX_LEN)
+  print('running an inference')
+  print(
+      model.forward(
+          idx,
+          input_pos,
+      )
+  )
+
+  # Convert model to tflite.
+  print('converting model to tflite')
+  edge_model = ai_edge_torch.convert(
+      model,
+      (
+          idx,
+          input_pos,
+      ),
+  )
+  edge_model.export('/tmp/toy_model.tflite')
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  return lowertools.exported_program_to_mlir_text(ep)
+
+
+def convert_toy_model_with_kv_cache(_) -> None:
+  """Converts a toy model with kv cache to tflite."""
+  dump_mlir = False
+
+  config = toy_model_with_kv_cache.get_model_config()
+  model = toy_model_with_kv_cache.ToyModelWithKVCache(config)
+  model.eval()
+  print('running an inference')
+  kv = kv_utils.KVCache.from_model_config(config)
+
+  tokens, input_pos = toy_model_with_kv_cache.get_sample_prefill_inputs()
+  decode_token, decode_input_pos = (
+      toy_model_with_kv_cache.get_sample_decode_inputs()
+  )
+  print(model.forward(tokens, input_pos, kv))
+
+  if dump_mlir:
+    mlir_text = _export_stablehlo_mlir(model, (tokens, input_pos, kv))
+    with open('/tmp/toy_model_with_external_kv.stablehlo.mlir', 'w') as f:
+      f.write(mlir_text)
+
+  # Convert model to tflite with 2 signatures (prefill + decode).
+  print('converting toy model to tflite with 2 signatures (prefill + decode)')
+  edge_model = (
+      ai_edge_torch.signature(
+          'prefill',
+          model,
+          sample_kwargs={
+              'tokens': tokens,
+              'input_pos': input_pos,
+              'kv_cache': kv,
+          },
+      )
+      .signature(
+          'decode',
+          model,
+          sample_kwargs={
+              'tokens': decode_token,
+              'input_pos': decode_input_pos,
+              'kv_cache': kv,
+          },
+      )
+      .convert()
+  )
+  edge_model.export('/tmp/toy_external_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+  app.run(convert_toy_model)

ai_edge_torch/generative/examples/test_models/toy_model.py

@@ -0,0 +1,156 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has a single-layer transformer block.
+from typing import Optional, Tuple
+
+from ai_edge_torch.generative.layers import builder
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+import torch
+from torch import nn
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+KV_CACHE_MAX_LEN = 100
+
+
+class ToySingleLayerModel(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_block = TransformerBlock(config.block_config(0), config)
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    # Toy model has only one block config.
+    attn_config = config.block_config(0).attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len,
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      mask: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.max_seq_len]
+
+    x = self.transformer_block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    return self.lm_head(x)
+
+
+class ToySingleLayerModelWeightSharing(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.lm_head = nn.Linear(
+        config.embedding_dim,
+        config.vocab_size,
+        bias=config.lm_head_use_bias,
+    )
+    self.lm_head.weight.data = self.tok_embedding.weight.data
+    self.transformer_block = TransformerBlock(config.block_config(0), config)
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    # Toy model has only one block config.
+    attn_config = config.block_config(0).attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len,
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      mask: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+
+    x = self.transformer_block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    res = self.lm_head(x)
+    return res
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32,
+      head_dim=4,
+      num_query_groups=4,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+      enable_kv_cache=False,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=400,
+      num_layers=1,
+      max_seq_len=KV_CACHE_MAX_LEN,
+      embedding_dim=128,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+  )
+  return config

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -0,0 +1,138 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""A toy example which has basic transformer block (w/ externalized KV-Cache)."""
+
+from typing import Optional, Tuple
+
+from ai_edge_torch.generative.layers import attention
+from ai_edge_torch.generative.layers import builder
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities.model_builder import ExportConfig
+import torch
+from torch import nn
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithKVCache(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    # Toy model has only one block config.
+    block_config = config.block_config(0)
+    self.transformer_blocks = nn.ModuleList(
+        attention.TransformerBlock(block_config, config)
+        for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    attn_config = block_config.attn_config
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(attn_config.rotary_percentage * attn_config.head_dim),
+        base=attn_config.rotary_base,
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len,
+    )
+    self.config = config
+
+  def forward(
+      self,
+      tokens: torch.Tensor,
+      input_pos: torch.Tensor,
+      kv_cache: kv_utils.KVCache,
+      mask: Optional[torch.Tensor] = None,
+      export_config: Optional[ExportConfig] = None,
+  ) -> Tuple[torch.Tensor, kv_utils.KVCache]:
+    x = self.tok_embedding(tokens)
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    if mask is None:
+      mask = self.mask_cache.index_select(2, input_pos)
+      mask = mask[:, :, :, : self.config.max_seq_len]
+
+    updated_kv_entries = []
+    for i, block in enumerate(self.transformer_blocks):
+      kv_entry = kv_cache.caches[i] if kv_cache else None
+      x, kv_entry = block(x, (cos, sin), mask, input_pos, kv_entry)
+      if kv_entry:
+        updated_kv_entries.append(kv_entry)
+
+    updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))
+
+    if export_config is not None:
+      if (
+          torch.numel(input_pos) > 1
+          and not export_config.output_logits_on_prefill
+      ):
+        return {'kv_cache': updated_kv_cache}
+
+    x = self.final_norm(x)
+    return {'logits': self.lm_head(x), 'kv_cache': updated_kv_cache}
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32,
+      head_dim=4,
+      num_query_groups=4,
+      rotary_base=10000,
+      rotary_percentage=1.0,
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  block_config = cfg.TransformerBlockConfig(
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      post_attention_norm_config=norm_config,
+  )
+  config = cfg.ModelConfig(
+      vocab_size=150,
+      num_layers=2,
+      max_seq_len=100,
+      embedding_dim=128,
+      block_configs=block_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  tokens = torch.unsqueeze(torch.arange(0, 100, dtype=torch.int), 0)
+  input_pos = torch.arange(0, 100, dtype=torch.int)
+  return tokens, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  tokens = torch.tensor([[1]], dtype=torch.int)
+  input_pos = torch.tensor([10])
+  return tokens, input_pos

ai_edge_torch/generative/examples/tiny_llama/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================