ai-edge-torch-nightly 0.5.0.dev20250409__py3-none-any.whl → 0.5.0.dev20250411__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -40,8 +40,8 @@ def _run_convert_passes(
       fx_passes.OptimizeLayoutTransposesPass(),
       fx_passes.CanonicalizePass(),
       fx_passes.BuildAtenCompositePass(),
-      fx_passes.CanonicalizePass(),
       fx_passes.RemoveNonUserOutputsPass(),
+      fx_passes.CastInputsBf16ToF32Pass(),
       fx_passes.CanonicalizePass(),
   ]
 

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -17,6 +17,7 @@ from typing import Sequence, Union
 
 from ai_edge_torch._convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass
 from ai_edge_torch._convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass
+from ai_edge_torch._convert.fx_passes.cast_inputs_bf16_to_f32_pass import CastInputsBf16ToF32Pass
 from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass
 from ai_edge_torch._convert.fx_passes.remove_non_user_outputs_pass import RemoveNonUserOutputsPass

ai_edge_torch/_convert/fx_passes/cast_inputs_bf16_to_f32_pass.py

@@ -0,0 +1,50 @@
+# 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.
+# ==============================================================================
+"""Pass to cast all inputs with torch.bfloat16 type to torch.float32."""
+
+
+from ai_edge_torch import fx_infra
+import torch
+
+
+def cast_f32(x):
+  return x.to(torch.float32)
+
+
+class CastInputsBf16ToF32Pass(fx_infra.ExportedProgramPassBase):
+  """This pass casts all inputs with torch.bfloat16 type to torch.float32."""
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    modified = False
+    for node in exported_program.graph.nodes:
+      if (
+          node.op == "placeholder"
+          and node.meta.get("val").dtype == torch.bfloat16
+      ):
+        if not node.users:
+          continue
+
+        modified = True
+        user = next(iter(node.users))
+        with exported_program.graph.inserting_before(user):
+          cast_node = exported_program.graph.call_function(
+              cast_f32,
+              (node,),
+          )
+          node.replace_all_uses_with(cast_node)
+          cast_node.replace_input_with(cast_node, node)
+
+    exported_program.graph_module.recompile()
+    return fx_infra.ExportedProgramPassResult(exported_program, modified)

ai_edge_torch/_convert/test/test_convert.py

@@ -553,6 +553,27 @@ class TestConvert(googletest.TestCase):
       self.fail(f"PT2E conversion failed: {err}")
     # pylint: enable=broad-except
 
+  def test_convert_model_with_bfloat16_inputs(self):
+    """Test converting a simple model with torch.bfloat16 input.
+
+    bf16 inputs would remain in converted model signature but be casted to f32
+    right after the model inputs.
+    """
+
+    class SampleModel(nn.Module):
+
+      def forward(self, x: torch.Tensor):
+        return (x + 1) * 1.2
+
+    model = SampleModel().eval()
+    args = (torch.randn(10, 10).to(torch.bfloat16),)
+    # pylint: disable=broad-except
+    try:
+      ai_edge_torch.convert(model, args)
+    except Exception as err:
+      self.fail(f"Conversion failed with bloat16 inputs: {err}")
+    # pylint: enable=broad-except
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_torch/generative/layers/experimental/attention.py

@@ -24,8 +24,7 @@ from typing import Optional, Tuple, Union
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
-from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
-from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
+from ai_edge_torch.generative.layers import sdpa_with_kv_update
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 import torch
@@ -147,7 +146,6 @@ class CausalSelfAttention(nn.Module):
     self.key_norm = builder.build_norm(config.head_dim, config.key_norm_config)
     self.config = config
     self.enable_hlfb = enable_hlfb
-    self.sdpa_func = sdpa.scaled_dot_product_attention
 
   def forward(
       self,
@@ -221,36 +219,8 @@ class CausalSelfAttention(nn.Module):
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 
-    # Transpose k/v to specific layout for GPU implementation.
-    b, _, n, h = q.shape
-    g = n // self.config.num_query_groups
-    # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
-    q = q.permute(0, 2, 1, 3).reshape(
-        1, b * self.config.num_query_groups, g * T, h
-    )
-
-    k = k.permute(0, 2, 1, 3).reshape(
-        1, -1, T, self.config.head_dim
-    )  # 1, bk, s, h
-    v = v.permute(0, 2, 3, 1).reshape(
-        1, -1, self.config.head_dim, T
-    )  # 1, bk, h, s
-
-    if kv_cache is not None:
-      kv_cache = kv_utils_experimental.update(kv_cache, input_pos, k, v)
-      k, v = kv_cache.k_cache, kv_cache.v_cache
-
-    sdpa_out = self.sdpa_func(
-        kv_cache,
-        q,
-        k,
-        v,
-        self.config.head_dim,
-        mask=mask,
-        softcap=self.config.logit_softcap,
-    )  # 1, bk, gt, h
-    sdpa_out = (
-        sdpa_out.reshape(B, -1, T, h).permute(0, 2, 1, 3).reshape(B, T, -1)
+    sdpa_out, kv_cache = sdpa_with_kv_update.sdpa_with_kv_update(
+        q, k, v, kv_cache, input_pos, mask, self.config
     )
 
     # Compute the output projection.

ai_edge_torch/generative/layers/experimental/kv_cache.py

@@ -44,7 +44,8 @@ def update(
   assert (
       cache.kv_layout == kv_utils.KV_LAYOUT_TRANSPOSED
   ), "KV entry must have transposed layout."
-  return _update_kv_impl_transposed(cache, input_pos, k_slice, v_slice)
+  update_kv_cache = _update_kv_impl_transposed
+  return update_kv_cache(cache, input_pos, k_slice, v_slice)
 
 
 def _get_slice_indices(

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -20,7 +20,7 @@ from typing import Optional
 
 from ai_edge_torch.generative.custom_ops import bmm_4d as bmm_lib
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import types
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 from multipledispatch import dispatch
 import torch
@@ -82,7 +82,6 @@ def _sdpa(k_type, v_type, *args, **kwargs):
   padded_logits = logits + mask
   padded_logits = padded_logits.reshape(1, bk, gt, s)
   probs = F.softmax(padded_logits, dim=-1).type_as(key)
-
   encoded = bmm_lib.bmm_4d(probs, value)
 
   return encoded  # 1, bk, gt, h

ai_edge_torch/generative/layers/kv_cache.py

@@ -20,7 +20,7 @@ from typing import Any, List, Tuple
 
 from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
 from ai_edge_torch.generative.layers import model_config
-from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import types
 import torch
 import torch.utils._pytree as pytree
 

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -0,0 +1,124 @@
+# 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.
+# ==============================================================================
+# Common utility functions for data loading etc.
+from dataclasses import dataclass
+from typing import Tuple
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
+from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa_default
+from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
+from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import types
+from multipledispatch import dispatch
+import torch
+
+
+def sdpa_with_kv_update(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  return sdpa_with_kv_update_impl(
+      kv.kv_layout[0](),  # key layout
+      kv.kv_layout[1](),  # value layout
+      query=query,
+      key=key,
+      value=value,
+      kv=kv,
+      input_pos=input_pos,
+      mask=mask,
+      config=config,
+  )
+
+
+@dispatch(types.BNTH, types.BNHT)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  kv = kwargs["kv"]
+  input_pos = kwargs["input_pos"]
+  mask = kwargs["mask"]
+  config = kwargs["config"]
+
+  # Transpose k/v to specific layout for GPU implementation.
+  b, seq_len, n, h = query.shape
+  g = n // config.num_query_groups
+  # btnh -> bnth -> b(kg)th -> 1(bk)(gt)h
+  query = query.permute(0, 2, 1, 3).reshape(
+      1, b * config.num_query_groups, g * seq_len, h
+  )
+
+  key = key.permute(0, 2, 1, 3).reshape(
+      1, -1, seq_len, config.head_dim
+  )  # 1, bk, s, h
+  value = value.permute(0, 2, 3, 1).reshape(
+      1, -1, config.head_dim, seq_len
+  )  # 1, bk, h, s
+
+  if kv is not None:
+    kv = kv_utils_experimental.update(kv, input_pos, key, value)
+    key, value = kv.k_cache, kv.v_cache
+
+  sdpa_out = sdpa.scaled_dot_product_attention(
+      kv,
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )  # 1, bk, gt, h
+  sdpa_out = (
+      sdpa_out.reshape(b, -1, seq_len, h)
+      .permute(0, 2, 1, 3)
+      .reshape(b, seq_len, -1)
+  )
+  return sdpa_out, kv
+
+
+@dispatch(object, object)
+def sdpa_with_kv_update_impl(
+    k_type, v_type, *args, **kwargs
+) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
+  query = kwargs["query"]
+  key = kwargs["key"]
+  value = kwargs["value"]
+  kv = kwargs["kv"]
+  input_pos = kwargs["input_pos"]
+  mask = kwargs["mask"]
+  config = kwargs["config"]
+
+  b, seq_len, _, _ = query.shape
+  if kv is not None:
+    kv = kv_utils.update(kv, input_pos, key, value)
+    key, value = kv.k_cache, kv.v_cache
+
+  sdpa_out = sdpa_default.scaled_dot_product_attention(
+      query,
+      key,
+      value,
+      config.head_dim,
+      mask=mask,
+      softcap=config.logit_softcap,
+  )
+  sdpa_out = sdpa_out.reshape(b, seq_len, -1)
+  return sdpa_out, kv

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -52,6 +52,7 @@ def torch_dtype_to_tf(dtype):
       torch.int32: tf.int32,
       torch.int16: tf.int16,
       torch.bool: tf.bool,
+      torch.bfloat16: tf.bfloat16,
   }.get(dtype)
 
 

ai_edge_torch/odml_torch/lowerings/_basic.py

@@ -301,3 +301,22 @@ def _aten_slice_scatter(lctx, self, src, dim=0, start=None, end=None, step=1):
   )
   out = stablehlo.select(pred, self, src)
   return out
+
+
+# Schema:
+#   - aten::_to_copy(Tensor self, *, ScalarType? dtype=None,
+#       Layout? layout=None, Device? device=None, bool? pin_memory=None,
+#       bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor
+@lower(torch.ops.aten._to_copy.default)
+def _aten_to_copy(
+    lctx, x: ir.Value, dtype: torch.dtype | None = None, **kwargs
+):
+  if not dtype:
+    return x
+
+  return stablehlo.convert(
+      ir.RankedTensorType.get(
+          x.type.shape, utils.torch_dtype_to_ir_element_type(dtype)
+      ),
+      x,
+  )

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -74,7 +74,6 @@ lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit)
 lower_by_torch_xla2(torch.ops.aten._native_batch_norm_legit_no_training)
 lower_by_torch_xla2(torch.ops.aten._pdist_forward)
 lower_by_torch_xla2(torch.ops.aten._softmax)
-lower_by_torch_xla2(torch.ops.aten._to_copy)
 lower_by_torch_xla2(torch.ops.aten._unsafe_index)
 lower_by_torch_xla2(torch.ops.aten._unsafe_view)
 lower_by_torch_xla2(torch.ops.aten.acos)

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -37,6 +37,7 @@ def torch_dtype_to_ir_element_type(dtype) -> ir.Type:
       torch.int16: functools.partial(ir.IntegerType.get_signless, 16),
       torch.int8: functools.partial(ir.IntegerType.get_signless, 8),
       torch.bool: functools.partial(ir.IntegerType.get_signless, 1),
+      torch.bfloat16: ir.BF16Type.get,
   }[dtype]
   return ty_get()
 

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250409.dist-info → ai_edge_torch_nightly-0.5.0.dev20250411.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250409
+Version: 0.5.0.dev20250411
 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.dev20250409.dist-info → ai_edge_torch_nightly-0.5.0.dev20250411.dist-info}/RECORD

@@ -2,16 +2,17 @@ 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=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=DEYqmCDZNmwuMxnxrFvcTEaDp6Z_BVHJaZMYjVQ2ijU,706
+ai_edge_torch/version.py,sha256=eSFSbkcpm9gvoLSoHH1AACS_xhGjYFxHgSKTrIlbpsQ,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
+ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
-ai_edge_torch/_convert/fx_passes/__init__.py,sha256=dG4WIICk0FqCH9euvbYHHsybRN7B1cYcuxN_OYxmjWo,1263
+ai_edge_torch/_convert/fx_passes/__init__.py,sha256=6LtGzzqT2IXprfI_vPYKhE7IuN5XmPG0xy-v0UtZ9yk,1361
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=a1KhqLetFb_efRHjX4T-zH0vF-U37Ha5I1CPIAsIluE,9211
 ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=3JyjiHpn17Zhfq3yGQXK5LMH71DQPXHb_4GOkP9uAjY,4251
+ai_edge_torch/_convert/fx_passes/cast_inputs_bf16_to_f32_pass.py,sha256=90YxLVAAkiA3qKr4Um__JmPeC1bTeA2PxBCj0GETq1Q,1748
 ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=Z6E3U7SYZvMl3Ivpqa3burVOLKFndEZuNmWKNxjq2mM,2386
 ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=HCOkj0k3NhaYbtfjE8HDXVmYhZ9fL5V_u6VunVh9mN4,2116
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=UKC-wM93-oe8spxyFqgybJ0TwnSRw8f-SOA2glCh2FA,890
@@ -26,7 +27,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=o6tuJkD-ESaQxLxJpN104qpchm3LCtPmHinzQxe6PSg,17226
+ai_edge_torch/_convert/test/test_convert.py,sha256=6vQa0UJn2L3qxR967_-vkfLrO7JdrLLBk4BfguOtHRI,17874
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -153,17 +154,17 @@ ai_edge_torch/generative/layers/attention.py,sha256=wLZ1jgUlcODBWgK3hnnhclHuuQDq
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 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/kv_cache.py,sha256=9kkFpB9msgUDStFxEyQYYsavKPP4Dgqb_NFcd4hA4aU,8502
+ai_edge_torch/generative/layers/kv_cache.py,sha256=WNH_Ab29eXKXs8HAm3Wmdv_LBzO6PQW5d34Eo6Yzgd0,8492
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=nLXvTkDAIHJQ0PTaWODF8oxJQoJ-K8D10cKR9229SAw,8355
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=oo9h7pi0GcuylRgp2yUuvUJCrhj03aoWt_fP7EDP4LM,3775
 ai_edge_torch/generative/layers/experimental/__init__.py,sha256=nz-K0h8DfiATHzR6s1_bCw2akUmHWffU1bDRSkIzSqI,592
-ai_edge_torch/generative/layers/experimental/attention.py,sha256=oW8cxv0pXcesnyGz6bXacRmlvHPfKNnJnls_Qb4L_aQ,8968
-ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=PlgL2bNNKasu3wFr3Iu9wbATWluWZt3_s4tzglJu2tM,2942
-ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=-ztTIgdec5gXkOVe6FXk3PMeS2HoL6-mBfDBdjQIcLQ,2808
-ai_edge_torch/generative/layers/experimental/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
+ai_edge_torch/generative/layers/experimental/attention.py,sha256=XYbo1KlmiMEuwArye0Ul86jEsdxLr1RG-usRpidZiT8,8001
+ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=zgpFVftOfllvjh9-UEBSvUbm152SnQETn29rUMMMvAM,2978
+ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=YFcIGOkaNb-vvQKjI-G9-bC2Z1W0O_qRyIZPlsLl72U,2797
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -193,6 +194,7 @@ ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
+ai_edge_torch/generative/utilities/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
 ai_edge_torch/generative/utilities/verifier.py,sha256=gsBGv-WgeP73Bmao29CiTIQOC1YZ43IUJcGzytpcZyM,12095
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=JsVmYrM_JEuN_smMHXUsRlo3Liapp7UyktbPpPARwDk,5386
@@ -203,7 +205,7 @@ ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=5kmOJWCc7sU1Hrqr1y17BtShUrss
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
 ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5zo3MlX1QYW4c,4513
-ai_edge_torch/lowertools/odml_torch_utils.py,sha256=dxg2pBuVhSZeY2Ouc0F6nHiZilWZmpSPA7I8kGqSkVI,8282
+ai_edge_torch/lowertools/odml_torch_utils.py,sha256=QRuS7S5lULRWEh3J1sWIsnKh-rbX7rd9tt6JJHbMPfo,8317
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
@@ -223,17 +225,17 @@ ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNi
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
 ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
-ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=Jq8_yAxC7ilzd6tOaRyBsOUEeenFF_EAC5haacZT4Pg,10247
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=4syWstepGiw3IKa8O7lciXywY7RFJ7OCWFMU1Lg3h-s,10777
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=VhmeGFnB5hrUsALiVWV96JJOqPDrTIWouHjTvLuT5eU,2477
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=h-YHW7xmvt9dpea-7Zj82HW7h5TKzW6GBEE13dIJQ40,11518
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=JRGLXW8EQ1L-vdiVTkD1kb4AnTU05eRwZ7Ke010hZmg,11473
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
-ai_edge_torch/odml_torch/lowerings/utils.py,sha256=-TzK1igPgR38oZkU1iPh-DZhlKVwuBtGWVC-y81PXzY,8935
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=uJaFbbgvYMI4-VFpFcMpaObNfBQl6nV0x8Yo8LaSAOE,8974
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
@@ -243,8 +245,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.dev20250409.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250409.dist-info/METADATA,sha256=kZwo6E79HLuM7_4E-Yw9erTzOnAAzio3Vy45hXNiC48,2051
-ai_edge_torch_nightly-0.5.0.dev20250409.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250409.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250409.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250411.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250411.dist-info/METADATA,sha256=NcEM89nSJwDWPZTuq137MbNXbp3zMHbtDZpR0_Q9z08,2051
+ai_edge_torch_nightly-0.5.0.dev20250411.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250411.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250411.dist-info/RECORD,,