ai-edge-torch-nightly 0.5.0.dev20250429__py3-none-any.whl → 0.5.0.dev20250430__py3-none-any.whl

ai_edge_torch/generative/layers/feed_forward_test.py

@@ -0,0 +1,54 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+from ai_edge_torch.generative.layers import feed_forward
+import torch
+import torch.nn.functional as F
+from absl.testing import absltest as googletest
+
+
+class FeedForwardTest(googletest.TestCase):
+
+  def test_sequential_feed_forward(self):
+    ff = feed_forward.SequentialFeedForward(
+        dim=10,
+        hidden_dim=10,
+        activation=F.silu,
+        use_bias=True,
+        use_glu=False,
+        pre_ff_norm=torch.nn.Identity(),
+        post_ff_norm=torch.nn.Identity(),
+    )
+    x = torch.ones((1, 10))
+    out = ff(x)
+    self.assertEqual(out.shape, (1, 10))
+
+  def test_gated_feed_forward(self):
+    ff = feed_forward.GatedFeedForward(
+        dim=10,
+        hidden_dim=10,
+        activation=F.silu,
+        use_bias=True,
+        use_glu=False,
+        pre_ff_norm=torch.nn.Identity(),
+        post_ff_norm=torch.nn.Identity(),
+    )
+    x = torch.ones((1, 10))
+    out = ff(x)
+    self.assertEqual(out.shape, (1, 10))
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250429"
+__version__ = "0.5.0.dev20250430"

{ai_edge_torch_nightly-0.5.0.dev20250429.dist-info → ai_edge_torch_nightly-0.5.0.dev20250430.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250429
+Version: 0.5.0.dev20250430
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.5.0.dev20250429.dist-info → ai_edge_torch_nightly-0.5.0.dev20250430.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=I820JmIf90_QKTKyhmQGVjX9U-WMGUVEo9_N-Q_aQuk,706
+ai_edge_torch/version.py,sha256=NOoOWF4AVnBmA1QQLYgLUaZYRPXTIKyH1nQ0DyJXftE,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -159,6 +159,7 @@ ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5
 ai_edge_torch/generative/layers/attention_utils_test.py,sha256=22gQ1gcRPkwqFG3_p82GZfRKVE3udEssSy58wNOqv0w,2431
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
+ai_edge_torch/generative/layers/feed_forward_test.py,sha256=8ZGy79BBpsyS6yKKDEKrDt249G5Mz-8VKWW7_WHx0u4,1655
 ai_edge_torch/generative/layers/kv_cache.py,sha256=b-7shzDaKexmvQF7P3SiAmIz4ZofjYWv3m5u71GojsA,10460
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=nLXvTkDAIHJQ0PTaWODF8oxJQoJ-K8D10cKR9229SAw,8355
@@ -246,8 +247,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/METADATA,sha256=05nMBPcVBVJcZhDI9SzsjryW3d4vpeeH_9H07RaA-PI,2051
-ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250429.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250430.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250430.dist-info/METADATA,sha256=QKdoe7QLXXeb8NfegbktiB0OuYjIbfDJjDIHPJ0B3sQ,2051
+ai_edge_torch_nightly-0.5.0.dev20250430.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250430.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250430.dist-info/RECORD,,