ai-edge-torch-nightly 0.6.0.dev20250522__py3-none-any.whl → 0.6.0.dev20250524__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -15,6 +15,7 @@
15
15
 
16
16
  """Common building blocks for Attention layer."""
17
17
 
18
+ import abc
18
19
  from typing import Optional, Tuple, Union
19
20
 
20
21
  from ai_edge_torch.generative.layers import builder
@@ -111,7 +112,42 @@ class TransformerBlock(nn.Module):
111
112
  return output if kv is None else (output, kv)
112
113
 
113
114
 
114
- class CausalSelfAttention(nn.Module):
115
+ class CausalSelfAttentionBase(nn.Module):
116
+ """Base class for causal self attention layer."""
117
+
118
+ def __init__(
119
+ self, dim: int, config: cfg.AttentionConfig, enable_hlfb: bool
120
+ ) -> None:
121
+ super().__init__()
122
+ self.dim = dim
123
+ self.config = config
124
+ self.enable_hlfb = enable_hlfb
125
+
126
+ self.query_norm = builder.build_norm(
127
+ self.config.head_dim, self.config.query_norm_config
128
+ )
129
+ self.key_norm = builder.build_norm(
130
+ self.config.head_dim, self.config.key_norm_config
131
+ )
132
+ self.value_norm = builder.build_norm(
133
+ self.config.head_dim, self.config.value_norm_config
134
+ )
135
+
136
+ @abc.abstractmethod
137
+ def forward(
138
+ self,
139
+ x: torch.Tensor,
140
+ rope: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
141
+ mask: Optional[torch.Tensor] = None,
142
+ input_pos: Optional[torch.Tensor] = None,
143
+ kv_cache: Optional[kv_utils.KVCacheEntry] = None,
144
+ lora: Optional[lora_utils.LoRAEntry] = None,
145
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, kv_utils.KVCacheEntry]]:
146
+ raise NotImplementedError()
147
+
148
+
149
+ class CausalSelfAttention(CausalSelfAttentionBase):
150
+ """Causal self attention layer implementation."""
115
151
 
116
152
  def __init__(
117
153
  self,
@@ -126,7 +162,7 @@ class CausalSelfAttention(nn.Module):
126
162
  config (cfg.AttentionConfig): attention specific configurations.
127
163
  enable_hlfb (bool): whether hlfb is enabled or not.
128
164
  """
129
- super().__init__()
165
+ super().__init__(dim, config, enable_hlfb)
130
166
  self.kv_cache = None
131
167
  qkv_shape = (
132
168
  config.num_heads + 2 * config.num_query_groups
@@ -137,12 +173,6 @@ class CausalSelfAttention(nn.Module):
137
173
  self.output_projection = nn.Linear(
138
174
  output_shape, dim, bias=config.output_proj_use_bias
139
175
  )
140
- self.query_norm = builder.build_norm(
141
- config.head_dim, config.query_norm_config
142
- )
143
- self.key_norm = builder.build_norm(config.head_dim, config.key_norm_config)
144
- self.config = config
145
- self.enable_hlfb = enable_hlfb
146
176
 
147
177
  def forward(
148
178
  self,
@@ -204,6 +234,7 @@ class CausalSelfAttention(nn.Module):
204
234
 
205
235
  q = self.query_norm(q)
206
236
  k = self.key_norm(k)
237
+ v = self.value_norm(v)
207
238
 
208
239
  q = q.reshape(B, T, -1, self.config.head_dim)
209
240
  k = k.reshape(B, T, -1, self.config.head_dim)
@@ -15,9 +15,9 @@
15
15
  # Builder class for individual components.
16
16
  from typing import Callable
17
17
 
18
+ from ai_edge_torch.generative.layers import normalization
18
19
  import ai_edge_torch.generative.layers.feed_forward as feed_forward
19
20
  import ai_edge_torch.generative.layers.model_config as cfg
20
- import ai_edge_torch.generative.layers.normalization as normalization
21
21
  import torch
22
22
  from torch import nn
23
23
  import torch.nn.functional as F
@@ -74,6 +74,8 @@ def build_norm(dim: int, config: cfg.NormalizationConfig):
74
74
  dim,
75
75
  eps=config.epsilon,
76
76
  zero_centered_gamma=config.zero_centered,
77
+ with_scale=config.with_scale,
78
+ scale_shift=config.scale_shift,
77
79
  enable_hlfb=config.enable_hlfb,
78
80
  )
79
81
  elif config.type == cfg.NormalizationType.LAYER_NORM:
@@ -107,20 +109,13 @@ def build_ff(dim: int, config: cfg.FeedForwardConfig):
107
109
  else:
108
110
  raise ValueError("Unsupported feedforward type.")
109
111
 
110
- activation = get_activation(config.activation)
111
-
112
112
  pre_ff_norm = build_norm(dim, config.pre_ff_norm_config)
113
113
  post_ff_norm = build_norm(dim, config.post_ff_norm_config)
114
114
 
115
115
  return ff_module(
116
116
  dim=dim,
117
- hidden_dim=config.intermediate_size,
118
- activation=activation,
119
- use_bias=config.use_bias,
120
- use_glu=(
121
- config.activation.type == cfg.ActivationType.GE_GLU
122
- or config.activation.type == cfg.ActivationType.SILU_GLU
123
- ),
117
+ activation=get_activation(config.activation),
118
+ config=config,
124
119
  pre_ff_norm=pre_ff_norm,
125
120
  post_ff_norm=post_ff_norm,
126
121
  )
@@ -14,45 +14,69 @@
14
14
  # ==============================================================================
15
15
  # Common building blocks for FeedForward layers.
16
16
 
17
- from typing import Callable, Optional
17
+ import abc
18
+ from typing import Callable
18
19
 
20
+ import ai_edge_torch.generative.layers.model_config as cfg
19
21
  import torch
20
22
  from torch import nn
21
23
 
22
24
 
23
- class SequentialFeedForward(nn.Module):
25
+ class FeedForwardBase(nn.Module):
26
+ """Base class for feedforward layer."""
27
+
28
+ def __init__(
29
+ self,
30
+ dim: int,
31
+ activation: Callable[[torch.Tensor], torch.Tensor],
32
+ config: cfg.FeedForwardConfig,
33
+ pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
34
+ post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
35
+ ):
36
+ super().__init__()
37
+ self.dim = dim
38
+ self.act = activation
39
+ self.config = config
40
+ self.hidden_dim = config.intermediate_size
41
+ self.use_bias = config.use_bias
42
+ self.use_glu = (
43
+ config.activation.type == cfg.ActivationType.GE_GLU
44
+ or config.activation.type == cfg.ActivationType.SILU_GLU
45
+ )
46
+ self.pre_ff_norm = pre_ff_norm
47
+ self.post_ff_norm = post_ff_norm
48
+
49
+ @abc.abstractmethod
50
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
51
+ raise NotImplementedError()
52
+
53
+
54
+ class SequentialFeedForward(FeedForwardBase):
24
55
  """Vanilla sequential Feedforward with customizable activation."""
25
56
 
26
57
  def __init__(
27
58
  self,
28
59
  dim: int,
29
- hidden_dim: int,
30
60
  activation: Callable[[torch.Tensor], torch.Tensor],
31
- use_bias=False,
32
- use_glu=False,
33
- pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
34
- post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
61
+ config: cfg.FeedForwardConfig,
62
+ pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
63
+ post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
35
64
  ):
36
65
  """Init function for feedforward layer.
37
66
 
38
67
  Args:
39
68
  dim (int): embedding size.
40
- hidden_dim (int): hidden dim size of the feedforward layer.
41
69
  activation (Callable): activation function used in this block.
42
- use_bias (Boolean): whether to use bias. Default is false.
43
- use_glu (Boolean): whether to use glu in activation. Default is false.
44
- pre_ff_norm (Callable): pre feedforward norm. Default is None.
45
- post_ff_norm (Callable): post feedforward norm. Default is None.
70
+ config (cfg.FeedForwardConfig): feedforward layer configuration.
71
+ pre_ff_norm (Callable): pre feedforward norm. Default is identity.
72
+ post_ff_norm (Callable): post feedforward norm. Default is identity.
46
73
  """
47
- super().__init__()
48
- self.act = activation
49
- if use_glu:
50
- self.w1 = nn.Linear(dim, hidden_dim * 2, bias=use_bias)
74
+ super().__init__(dim, activation, config, pre_ff_norm, post_ff_norm)
75
+ if self.use_glu:
76
+ self.w1 = nn.Linear(dim, self.hidden_dim * 2, bias=self.use_bias)
51
77
  else:
52
- self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
53
- self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
54
- self.pre_ff_norm = pre_ff_norm if pre_ff_norm else lambda x: x
55
- self.post_ff_norm = post_ff_norm if post_ff_norm else lambda x: x
78
+ self.w1 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
79
+ self.w2 = nn.Linear(self.hidden_dim, dim, bias=self.use_bias)
56
80
 
57
81
  def forward(self, x):
58
82
  """Forward pass for Feedforward layer.
@@ -68,7 +92,7 @@ class SequentialFeedForward(nn.Module):
68
92
  return self.post_ff_norm(out)
69
93
 
70
94
 
71
- class GatedFeedForward(nn.Module):
95
+ class GatedFeedForward(FeedForwardBase):
72
96
  """Gated Feedforward with customizable activation.
73
97
 
74
98
  https://arxiv.org/pdf/2002.05202v1.pdf
@@ -77,34 +101,48 @@ class GatedFeedForward(nn.Module):
77
101
  def __init__(
78
102
  self,
79
103
  dim: int,
80
- hidden_dim: int,
81
104
  activation: Callable[[torch.Tensor], torch.Tensor],
82
- use_bias=False,
83
- use_glu=False,
84
- pre_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
85
- post_ff_norm: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
105
+ config: cfg.FeedForwardConfig,
106
+ pre_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
107
+ post_ff_norm: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
86
108
  ):
87
109
  """Init function for feedforward layer.
88
110
 
89
111
  Args:
90
112
  dim (int): embedding size.
91
- hidden_dim (int): hidden dim size of the feedforward layer.
92
113
  activation (Callable): activation function used in this block.
93
- use_bias (Boolean): whether to use bias. Default is false.
94
- use_glu (Boolean): whether to use glu in activation. Default is false.
95
- pre_ff_norm (Callable): pre feedforward norm. Default is None.
96
- post_ff_norm (Callable): post feedforward norm. Default is None.
114
+ pre_ff_norm (Callable): pre feedforward norm. Default is identity.
115
+ post_ff_norm (Callable): post feedforward norm. Default is identity.
116
+ config (cfg.FeedForwardConfig): feedforward layer configuration.
97
117
  """
98
- super().__init__()
99
- self.act = activation
100
- if use_glu:
101
- self.w1 = nn.Linear(dim, hidden_dim * 2, bias=use_bias)
118
+ super().__init__(dim, activation, config, pre_ff_norm, post_ff_norm)
119
+
120
+ if self.use_glu:
121
+ assert (
122
+ self.config.use_separate_gating
123
+ ), 'use_separate_gating must be True for GE_GLU | SILU_GLU activation.'
124
+
125
+ if self.config.use_separate_gating:
126
+ if self.use_glu:
127
+ self.w1 = nn.Linear(dim, self.hidden_dim * 2, bias=self.use_bias)
128
+ else:
129
+ self.w1 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
130
+ self.w3 = nn.Linear(dim, self.hidden_dim, bias=self.use_bias)
102
131
  else:
103
- self.w1 = nn.Linear(dim, hidden_dim, bias=use_bias)
104
- self.w2 = nn.Linear(hidden_dim, dim, bias=use_bias)
105
- self.w3 = nn.Linear(dim, hidden_dim, bias=use_bias)
106
- self.pre_ff_norm = pre_ff_norm if pre_ff_norm else lambda x: x
107
- self.post_ff_norm = post_ff_norm if post_ff_norm else lambda x: x
132
+ self.w_gating = nn.Parameter(
133
+ torch.ones((2, dim, self.hidden_dim), dtype=torch.float32),
134
+ requires_grad=False,
135
+ )
136
+ self.gating_bias = (
137
+ nn.Parameter(
138
+ torch.zeros((2, self.hidden_dim), dtype=torch.float32),
139
+ requires_grad=False,
140
+ )
141
+ if self.use_bias
142
+ else torch.zeros((2, self.hidden_dim), dtype=torch.float32)
143
+ )
144
+
145
+ self.w2 = nn.Linear(self.hidden_dim, dim, bias=self.use_bias)
108
146
 
109
147
  def forward(self, x):
110
148
  """Forward pass for Feedforward layer.
@@ -116,5 +154,12 @@ class GatedFeedForward(nn.Module):
116
154
  torch.Tensor: output tensor after feedforward.
117
155
  """
118
156
  x_norm = self.pre_ff_norm(x)
119
- out = self.w2(self.act(self.w1(x_norm)) * self.w3(x_norm))
157
+ if self.config.use_separate_gating:
158
+ out = self.w2(self.act(self.w1(x_norm)) * self.w3(x_norm))
159
+ else:
160
+ out = self.w2(
161
+ self.act(torch.matmul(x_norm, self.w_gating[0]) + self.gating_bias[0])
162
+ * (torch.matmul(x_norm, self.w_gating[1]) + self.gating_bias[1])
163
+ )
164
+
120
165
  return self.post_ff_norm(out)
@@ -14,6 +14,7 @@
14
14
  # ==============================================================================
15
15
 
16
16
  from ai_edge_torch.generative.layers import feed_forward
17
+ from ai_edge_torch.generative.layers import model_config as cfg
17
18
  import torch
18
19
  import torch.nn.functional as F
19
20
  from absl.testing import absltest as googletest
@@ -22,28 +23,32 @@ from absl.testing import absltest as googletest
22
23
  class FeedForwardTest(googletest.TestCase):
23
24
 
24
25
  def test_sequential_feed_forward(self):
26
+ ff_config = cfg.FeedForwardConfig(
27
+ type=cfg.FeedForwardType.SEQUENTIAL,
28
+ activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
29
+ intermediate_size=10,
30
+ use_bias=True,
31
+ )
25
32
  ff = feed_forward.SequentialFeedForward(
26
33
  dim=10,
27
- hidden_dim=10,
28
34
  activation=F.silu,
29
- use_bias=True,
30
- use_glu=False,
31
- pre_ff_norm=torch.nn.Identity(),
32
- post_ff_norm=torch.nn.Identity(),
35
+ config=ff_config,
33
36
  )
34
37
  x = torch.ones((1, 10))
35
38
  out = ff(x)
36
39
  self.assertEqual(out.shape, (1, 10))
37
40
 
38
41
  def test_gated_feed_forward(self):
42
+ ff_config = cfg.FeedForwardConfig(
43
+ type=cfg.FeedForwardType.GATED,
44
+ activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
45
+ intermediate_size=10,
46
+ use_bias=True,
47
+ )
39
48
  ff = feed_forward.GatedFeedForward(
40
49
  dim=10,
41
- hidden_dim=10,
42
50
  activation=F.silu,
43
- use_bias=True,
44
- use_glu=False,
45
- pre_ff_norm=torch.nn.Identity(),
46
- post_ff_norm=torch.nn.Identity(),
51
+ config=ff_config,
47
52
  )
48
53
  x = torch.ones((1, 10))
49
54
  out = ff(x)
@@ -69,10 +69,32 @@ class NormalizationConfig:
69
69
  enable_hlfb: bool = True
70
70
  epsilon: float = 1e-5
71
71
  zero_centered: bool = False
72
+ # Whether to use a scale parameter in the normalization.
73
+ with_scale: bool = False
74
+ # The shift to apply to the scale parameter.
75
+ scale_shift: float = 0.0
72
76
  # Number of groups used in group normalization.
73
77
  group_num: Optional[float] = None
74
78
 
75
79
 
80
+ # Exprimental feature and may subject to change.
81
+ class KVCacheUpdateStrategy(enum.Enum):
82
+ """Different alignment strategies of the KV cache.
83
+
84
+ Due to restrictions from different devices, we may need to apply different
85
+ alignment strategies to the KV cache during Attention layer's cache update.
86
+
87
+ Available options:
88
+ INPLACE: Update the existing cache in place using indexes.
89
+ PREPEND_LEFT: Append the new kv to the left of the existing cache. When this
90
+ cache update is applied, the newer kvs will always be prepended at the
91
+ beginning of the cache.
92
+ """
93
+
94
+ INPLACE = enum.auto()
95
+ PREPEND_LEFT = enum.auto()
96
+
97
+
76
98
  @dataclasses.dataclass
77
99
  class AttentionConfig:
78
100
  """Attention model's parameters."""
@@ -108,6 +130,12 @@ class AttentionConfig:
108
130
  key_norm_config: NormalizationConfig = dataclasses.field(
109
131
  default_factory=NormalizationConfig
110
132
  )
133
+ # The normalization applied to value projection's output.
134
+ value_norm_config: NormalizationConfig = dataclasses.field(
135
+ default_factory=NormalizationConfig
136
+ )
137
+ # Whether the KV cache is shared with the previous attention block.
138
+ kv_shared: bool = False
111
139
  relative_attention_num_buckets: int = 0
112
140
  relative_attention_max_distance: int = 0
113
141
  # Softcap on the output logits.
@@ -118,6 +146,8 @@ class AttentionConfig:
118
146
  sliding_window_size: Optional[int] = None
119
147
  # The default causal mask value used by attention layer.
120
148
  causal_mask_value: float = float("-inf")
149
+ # The update strategy of the KV cache. Default to INPLACE.
150
+ kvcache_update_strategy: KVCacheUpdateStrategy = KVCacheUpdateStrategy.INPLACE
121
151
 
122
152
 
123
153
  @dataclasses.dataclass
@@ -135,6 +165,9 @@ class FeedForwardConfig:
135
165
  type: FeedForwardType
136
166
  activation: ActivationConfig
137
167
  intermediate_size: int
168
+ # Whether to use two separate gating parameters or a single one in
169
+ # GatedFeedForward.
170
+ use_separate_gating: bool = True
138
171
  use_bias: bool = False
139
172
  # The normalization applied to feed forward's input.
140
173
  pre_ff_norm_config: NormalizationConfig = dataclasses.field(
ai_edge_torch/version.py CHANGED
@@ -13,4 +13,4 @@
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
15
 
16
- __version__ = "0.6.0.dev20250522"
16
+ __version__ = "0.6.0.dev20250524"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ai-edge-torch-nightly
3
- Version: 0.6.0.dev20250522
3
+ Version: 0.6.0.dev20250524
4
4
  Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
5
5
  Home-page: https://github.com/google-ai-edge/ai-edge-torch
6
6
  Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
2
2
  ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
3
3
  ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
4
4
  ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
5
- ai_edge_torch/version.py,sha256=khTRYaBBm9J5iM-bJK8PgD2F4rAYy3w4cGEKRbQ498M,706
5
+ ai_edge_torch/version.py,sha256=YonvsAweNsFCQhDNTkzmc937o8WNdTfCG3AcozW5WZk,706
6
6
  ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
7
7
  ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
8
8
  ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -169,18 +169,18 @@ ai_edge_torch/generative/examples/tiny_llama/verify_util.py,sha256=_zYGqP4HO_Stc
169
169
  ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
170
170
  ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
171
171
  ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
172
- ai_edge_torch/generative/layers/attention.py,sha256=uK1ih2kxPZherwi-pGSm8B--NNWnQ8npEAfgcjMIkEY,12964
172
+ ai_edge_torch/generative/layers/attention.py,sha256=RaXENRRQo1MsLdt3U8h3kYTCmd6imHQ-aCXtmPXCh_o,13911
173
173
  ai_edge_torch/generative/layers/attention_test.py,sha256=9v8v96TLyFPdqxEylU1JOAeRFAp2s0YoDHZN83SFJJA,4764
174
174
  ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
175
175
  ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
176
- ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
176
+ ai_edge_torch/generative/layers/builder.py,sha256=2bUgkyowDkDznkF8XaHyZs4nowHr1QEHYLM7pMaFmIk,4921
177
177
  ai_edge_torch/generative/layers/einsum.py,sha256=EsZSWNVWUs0-1plp4TBnhP4ZhaRDBa2VlDO6hWpUAqU,1288
178
178
  ai_edge_torch/generative/layers/einsum_test.py,sha256=ltIE773bvvNLv_9aLQxFwe1MgQ762sez0c5E2tejxuA,1079
179
- ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
180
- ai_edge_torch/generative/layers/feed_forward_test.py,sha256=8ZGy79BBpsyS6yKKDEKrDt249G5Mz-8VKWW7_WHx0u4,1655
179
+ ai_edge_torch/generative/layers/feed_forward.py,sha256=_GmtHxwL068l9gh_F_WFcFk7La-Tl5SfoQ9v2hMabZM,5541
180
+ ai_edge_torch/generative/layers/feed_forward_test.py,sha256=Y5l1eC9NgfYixHcfIfE1W4FGh7oC-9UGGyHdKS9tQKc,1880
181
181
  ai_edge_torch/generative/layers/kv_cache.py,sha256=b-7shzDaKexmvQF7P3SiAmIz4ZofjYWv3m5u71GojsA,10460
182
182
  ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
183
- ai_edge_torch/generative/layers/model_config.py,sha256=H1MpjP1Ij1r4DEcE4cQ_6A8h0QvUjCkuGATXMkIMIWg,8570
183
+ ai_edge_torch/generative/layers/model_config.py,sha256=0FH3UJPVnEhgBO4eUlNaHuQBDo_OKH17ChG5-Ybj2T4,9895
184
184
  ai_edge_torch/generative/layers/normalization.py,sha256=ijwCpi22NLX-Sygwy5sK9l9WjGvbPIhZvVwoBAonWAo,7014
185
185
  ai_edge_torch/generative/layers/normalization_test.py,sha256=zwurZly-TgFxdgVVdpzu9vCpcLbd5RYt_gKg9Lfg1jI,2248
186
186
  ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
@@ -267,8 +267,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
267
267
  ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
268
268
  ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
269
269
  ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
270
- ai_edge_torch_nightly-0.6.0.dev20250522.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
271
- ai_edge_torch_nightly-0.6.0.dev20250522.dist-info/METADATA,sha256=4AO6S0O6AyT4GmWBjgWBjxRvbwD34vv4IoRR9kZW7V8,2074
272
- ai_edge_torch_nightly-0.6.0.dev20250522.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
273
- ai_edge_torch_nightly-0.6.0.dev20250522.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
274
- ai_edge_torch_nightly-0.6.0.dev20250522.dist-info/RECORD,,
270
+ ai_edge_torch_nightly-0.6.0.dev20250524.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
271
+ ai_edge_torch_nightly-0.6.0.dev20250524.dist-info/METADATA,sha256=ByeXgIhTWd4UpnG3N18UGahFi_yoFD75CBlUlR2qBO4,2074
272
+ ai_edge_torch_nightly-0.6.0.dev20250524.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
273
+ ai_edge_torch_nightly-0.6.0.dev20250524.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
274
+ ai_edge_torch_nightly-0.6.0.dev20250524.dist-info/RECORD,,