ai-edge-torch-nightly 0.2.0.dev20240606__py3-none-any.whl → 0.2.0.dev20240608__py3-none-any.whl

ai_edge_torch/generative/utilities/autoencoder_loader.py

@@ -0,0 +1,298 @@
+# 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 Dict, List, Tuple
+
+import torch
+
+import ai_edge_torch.generative.layers.model_config as layers_config
+import ai_edge_torch.generative.layers.unet.model_config as unet_config
+import ai_edge_torch.generative.utilities.loader as loader
+
+
+@dataclass
+class ResidualBlockTensorNames:
+  norm_1: str = None
+  conv_1: str = None
+  norm_2: str = None
+  conv_2: str = None
+  residual_layer: str = None
+
+
+@dataclass
+class AttnetionBlockTensorNames:
+  norm: str = None
+  fused_qkv_proj: str = None
+  output_proj: str = None
+
+
+@dataclass
+class MidBlockTensorNames:
+  residual_block_tensor_names: List[ResidualBlockTensorNames]
+  attention_block_tensor_names: List[AttnetionBlockTensorNames]
+
+
+@dataclass
+class UpDecoderBlockTensorNames:
+  residual_block_tensor_names: List[ResidualBlockTensorNames]
+  upsample_conv: str = None
+
+
+def _map_to_converted_state(
+    state: Dict[str, torch.Tensor],
+    state_param: str,
+    converted_state: Dict[str, torch.Tensor],
+    converted_state_param: str,
+):
+  converted_state[f"{converted_state_param}.weight"] = state.pop(
+      f"{state_param}.weight"
+  )
+  if f"{state_param}.bias" in state:
+    converted_state[f"{converted_state_param}.bias"] = state.pop(f"{state_param}.bias")
+
+
+class AutoEncoderModelLoader(loader.ModelLoader):
+
+  @dataclass
+  class TensorNames:
+    quant_conv: str = None
+    post_quant_conv: str = None
+    conv_in: str = None
+    conv_out: str = None
+    final_norm: str = None
+    mid_block_tensor_names: MidBlockTensorNames = None
+    up_decoder_blocks_tensor_names: List[UpDecoderBlockTensorNames] = None
+
+  def __init__(self, file_name: str, names: TensorNames):
+    """AutoEncoderModelLoader constructor. Can be used to load encoder and decoder models.
+
+    Args:
+        file_name (str): Path to the checkpoint. Can be a directory or an
+          exact file.
+        names (TensorNames): An instance of `TensorNames` to determine mappings.
+    """
+    self._file_name = file_name
+    self._names = names
+    self._loader = self._get_loader()
+
+  def load(
+      self, model: torch.nn.Module, strict: bool = True
+  ) -> Tuple[List[str], List[str]]:
+    """Load the model from the checkpoint.
+
+    Args:
+        model (torch.nn.Module): The pytorch model that needs to be loaded.
+        strict (bool, optional): Whether the converted keys are strictly
+          matched. Defaults to True.
+
+    Returns:
+        missing_keys (List[str]): a list of str containing the missing keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected keys.
+
+    Raises:
+        ValueError: If conversion results in unmapped tensors and strict mode is
+          enabled.
+    """
+    state = self._loader(self._file_name)
+    converted_state = dict()
+    if self._names.quant_conv is not None:
+      _map_to_converted_state(
+          state, self._names.quant_conv, converted_state, "quant_conv"
+      )
+    if self._names.post_quant_conv is not None:
+      _map_to_converted_state(
+          state, self._names.post_quant_conv, converted_state, "post_quant_conv"
+      )
+    if self._names.conv_in is not None:
+      _map_to_converted_state(state, self._names.conv_in, converted_state, "conv_in")
+    if self._names.conv_out is not None:
+      _map_to_converted_state(state, self._names.conv_out, converted_state, "conv_out")
+    if self._names.final_norm is not None:
+      _map_to_converted_state(
+          state, self._names.final_norm, converted_state, "final_norm"
+      )
+    self._map_mid_block(
+        state,
+        converted_state,
+        model.config.mid_block_config,
+        self._names.mid_block_tensor_names,
+    )
+
+    reversed_block_out_channels = list(reversed(model.config.block_out_channels))
+    block_out_channels = reversed_block_out_channels[0]
+    for i, out_channels in enumerate(reversed_block_out_channels):
+      prev_output_channel = block_out_channels
+      block_out_channels = out_channels
+      not_final_block = i < len(reversed_block_out_channels) - 1
+      self._map_up_decoder_block(
+          state,
+          converted_state,
+          f"up_decoder_blocks.{i}",
+          unet_config.UpDecoderBlock2DConfig(
+              in_channels=prev_output_channel,
+              out_channels=block_out_channels,
+              normalization_config=model.config.normalization_config,
+              activation_type=model.config.activation_type,
+              num_layers=model.config.layers_per_block,
+              add_upsample=not_final_block,
+              upsample_conv=True,
+          ),
+          self._names.up_decoder_blocks_tensor_names[i],
+      )
+    if strict and state:
+      raise ValueError(
+          f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
+      )
+    return model.load_state_dict(converted_state, strict=strict)
+
+  def _map_residual_block(
+      self,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      tensor_names: ResidualBlockTensorNames,
+      converted_state_param_prefix: str,
+      config: unet_config.ResidualBlock2DConfig,
+  ):
+    _map_to_converted_state(
+        state,
+        tensor_names.norm_1,
+        converted_state,
+        f"{converted_state_param_prefix}.norm_1",
+    )
+    _map_to_converted_state(
+        state,
+        tensor_names.conv_1,
+        converted_state,
+        f"{converted_state_param_prefix}.conv_1",
+    )
+    _map_to_converted_state(
+        state,
+        tensor_names.norm_2,
+        converted_state,
+        f"{converted_state_param_prefix}.norm_2",
+    )
+    _map_to_converted_state(
+        state,
+        tensor_names.conv_2,
+        converted_state,
+        f"{converted_state_param_prefix}.conv_2",
+    )
+    if config.in_channels != config.out_channels:
+      _map_to_converted_state(
+          state,
+          tensor_names.residual_layer,
+          converted_state,
+          f"{converted_state_param_prefix}.residual_layer",
+      )
+
+  def _map_attention_block(
+      self,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      tensor_names: AttnetionBlockTensorNames,
+      converted_state_param_prefix: str,
+      config: unet_config.AttentionBlock2DConfig,
+  ):
+    if config.normalization_config.type != layers_config.NormalizationType.NONE:
+      _map_to_converted_state(
+          state,
+          tensor_names.norm,
+          converted_state,
+          f"{converted_state_param_prefix}.norm",
+      )
+    attention_layer_prefix = f"{converted_state_param_prefix}.attention"
+    _map_to_converted_state(
+        state,
+        tensor_names.fused_qkv_proj,
+        converted_state,
+        f"{attention_layer_prefix}.qkv_projection",
+    )
+    _map_to_converted_state(
+        state,
+        tensor_names.output_proj,
+        converted_state,
+        f"{attention_layer_prefix}.output_projection",
+    )
+
+  def _map_mid_block(
+      self,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      config: unet_config.MidBlock2DConfig,
+      tensor_names: MidBlockTensorNames,
+  ):
+    converted_state_param_prefix = "mid_block"
+    residual_block_config = unet_config.ResidualBlock2DConfig(
+        in_channels=config.in_channels,
+        out_channels=config.in_channels,
+        time_embedding_channels=config.time_embedding_channels,
+        normalization_config=config.normalization_config,
+        activation_type=config.activation_type,
+    )
+    self._map_residual_block(
+        state,
+        converted_state,
+        tensor_names.residual_block_tensor_names[0],
+        f"{converted_state_param_prefix}.resnets.0",
+        residual_block_config,
+    )
+    for i in range(config.num_layers):
+      if config.attention_block_config:
+        self._map_attention_block(
+            state,
+            converted_state,
+            tensor_names.attention_block_tensor_names[i],
+            f"{converted_state_param_prefix}.attentions.{i}",
+            config.attention_block_config,
+        )
+      self._map_residual_block(
+          state,
+          converted_state,
+          tensor_names.residual_block_tensor_names[i + 1],
+          f"{converted_state_param_prefix}.resnets.{i+1}",
+          residual_block_config,
+      )
+
+  def _map_up_decoder_block(
+      self,
+      state: Dict[str, torch.Tensor],
+      converted_state: Dict[str, torch.Tensor],
+      converted_state_param_prefix: str,
+      config: unet_config.UpDecoderBlock2DConfig,
+      tensor_names: UpDecoderBlockTensorNames,
+  ):
+    for i in range(config.num_layers):
+      input_channels = config.in_channels if i == 0 else config.out_channels
+      self._map_residual_block(
+          state,
+          converted_state,
+          tensor_names.residual_block_tensor_names[i],
+          f"{converted_state_param_prefix}.resnets.{i}",
+          unet_config.ResidualBlock2DConfig(
+              in_channels=input_channels,
+              out_channels=config.out_channels,
+              time_embedding_channels=config.time_embedding_channels,
+              normalization_config=config.normalization_config,
+              activation_type=config.activation_type,
+          ),
+      )
+    if config.add_upsample and config.upsample_conv:
+      _map_to_converted_state(
+          state,
+          tensor_names.upsample_conv,
+          converted_state,
+          f"{converted_state_param_prefix}.upsample_conv",
+      )

ai_edge_torch/generative/utilities/loader.py

@@ -16,7 +16,7 @@
 from dataclasses import dataclass
 import glob
 import os
-from typing import Callable, Dict
+from typing import Callable, Dict, List, Tuple
 
 from safetensors import safe_open
 import torch
@@ -129,8 +129,10 @@ class ModelLoader:
     self._names = names
     self._loader = self._get_loader()
 
-  def load(self, model: torch.nn.Module, strict: bool = True):
-    """Load the model from the checkpoint
+  def load(
+      self, model: torch.nn.Module, strict: bool = True
+  ) -> Tuple[List[str], List[str]]:
+    """Load the model from the checkpoint.
 
     Args:
         model (torch.nn.Module): The pytorch model that needs to be loaded.
@@ -138,8 +140,8 @@ class ModelLoader:
           matched. Defaults to True.
 
     Returns:
-        missing_keys (List[str]): a list of str containing the missing keys
-        unexpected_keys (List[str]): a list of str containing the unexpected keys
+        missing_keys (List[str]): a list of str containing the missing keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected keys.
 
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is

{ai_edge_torch_nightly-0.2.0.dev20240606.dist-info → ai_edge_torch_nightly-0.2.0.dev20240608.dist-info}/METADATA

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

@@ -1,13 +1,13 @@
 ai_edge_torch/__init__.py,sha256=FPMmuFU3pyMREtjB_san1fy_0PFtAsgA0VZfOYvDrb4,1008
 ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/convert/conversion.py,sha256=OlqiB17iTXBxUWX0LX4OU8OmZzLP94STk66LR-QF_V8,4043
+ai_edge_torch/convert/conversion.py,sha256=GN2Js232u_5Y118wg3qIfEoYewxbxLl3TpSnO6osi8c,4029
 ai_edge_torch/convert/conversion_utils.py,sha256=NpVm3Ms81_cIW5IYgGsr0BVganJJgBKWVBDe5h_ZaGE,11021
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
-ai_edge_torch/convert/fx_passes/__init__.py,sha256=Ll2nNwufjcV5nSruQPXiloq7F1E7pWJ2T5clXmy1lk8,2825
+ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
 ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=2yqUwJJ2R233_X9FNMOP9oYRTTzH34TR_BIUj-wfnKw,7080
-ai_edge_torch/convert/fx_passes/build_upsample_bilinear2d_composite_pass.py,sha256=76XYoIlFDgrzp5QemoaEalPFcEbfszkEH_PLvO1ASCk,2607
+ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=jB27GlDC8x36nn35aiq97uKERiq4KXSUZ7tv7yc0Gl4,3223
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=VA9bekxPVhLk4MYlIRXnOzrSnbCtUmGj7OQ_fJcKQtc,795
@@ -41,10 +41,10 @@ ai_edge_torch/generative/examples/phi2/phi2.py,sha256=VvigzPQ_LJHeADTsMliwFwPe2B
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=yUCJemEh4n8ez-yLgVU0HZAki-PZ9nY04DFjgpx9PUc,3698
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=MI73RjOeD4Kh7AL0j5_QXiZq-rl_qCdibSE6eCQCyeY,3804
-ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=AgVAdUbSkHXONVUjAyBQEXhIUUlinf9kNljcBpWnj3A,3276
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=NmgDo5uAefrhMUbYku0TKHlqzO0NVWI_M1ue8tddQR4,4024
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=Z1bnZvYtPdwNy706kixVDfL32X-R87B_WF3CcHwiz0o,11038
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=TfbfsmuKoGsBENF9fYIAN_SMEQNhj-kjNdqQXFJGxpg,7784
-ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=L6hLaMQGb8-_BwSvTLIuDnZwfTqn0K4swBUjfPnYWZo,2341
+ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=mgbxkeFDMkNIGmnbcFTIFPu8EWKokghiviYIOB2lE3Q,3437
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=FCbnwlkpYYb-tF7KscbSYjNEdg7XnuLju1cDuIRoQv8,8277
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=r9RqbyNvuvXOGu3ojtl7ZmbC7o4Pt8aUKAhN1yCdtEc,3397
 ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=NFpOfA4KN0JpShm5QvuYbQYZ844NzexWD8nV3WjMOZM,2397
@@ -65,15 +65,19 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TI
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=hVGpuI8gpj4Rn9k4otsRE22MSLFHBDlUOgioY6Ru6VI,5629
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=zNIBXxCOA5Mz_F_dfBbKpIovhtcB6q5a-i8oAxls1d0,7071
+ai_edge_torch/generative/layers/attention.py,sha256=Z8gXHYs6h8gaRiYAdvYUbHzg_2EmqfxiChsf_SYraAc,7902
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
-ai_edge_torch/generative/layers/builder.py,sha256=WLTeDId9t3Xwt0h1zxzqoYyFvfrNzPKLskcl39q8Aqw,3403
+ai_edge_torch/generative/layers/builder.py,sha256=8cPL1NAutjT6Dwtyy2X7NSaTl9WCUJM5SIrBIDcEvVY,3520
 ai_edge_torch/generative/layers/feed_forward.py,sha256=4j2QaSCw59Jkk_ixKDpKEj7FLRauzuExTiSNRzAjAhE,2820
 ai_edge_torch/generative/layers/kv_cache.py,sha256=4uiZLO3om5G3--kT04Jt0esEYznbkJ7QLzSHfb8mjc4,3090
-ai_edge_torch/generative/layers/model_config.py,sha256=2zT9nyoyuuyk5ziiww0VSJ6_JO7pDf7uOYbO9O3OQc4,4249
+ai_edge_torch/generative/layers/model_config.py,sha256=72DXOsFC0buvzZp6YyVjuTVrpphAubBJ5NJWfs3kEwk,4362
 ai_edge_torch/generative/layers/normalization.py,sha256=M27eW3TcNK20oaXClXtfnu0lLWrAGrSKSsbegRWnj3c,1867
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=12SsCuoRuLNCwnFGe_pHDOZEBwBcqXs87Aj0PaWWw4E,1383
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=dYafGC205QE5CLIbBTCI-7eVvEGZEHzs1toPEhemeDs,3391
+ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=KuZd2oZhkCQSknSgXMBla-sfYBPUv5bZNf9RYKXHfGg,10052
+ai_edge_torch/generative/layers/unet/builder.py,sha256=iH0_nuY9TF2ap5h1JbGNCOonPTfrXQHcF8U0slrIREM,1210
+ai_edge_torch/generative/layers/unet/model_config.py,sha256=LeRGB34fQ73UlknlFpjM9U-SZIRcQDnSmDltJivX-UA,4044
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=t-YwyKSPAG-OZC1DfH-0vfie2RHHpTSQjxUY-tmhu5g,1543
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=ffBALrrbrfiG_mrOr-f3B1Gc6PlAma9gtvVnfP7SDzI,1862
@@ -86,7 +90,8 @@ ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-y
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=i_SAW-hD8SaHuopMZI9IuXXDFn5uSTJa1nKZhaC3dAQ,6811
 ai_edge_torch/generative/test/test_quantize.py,sha256=f70sH1ZFzdCwYj0MG-eg54WOC4LasR0D8CTUYpjxZYM,3728
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/loader.py,sha256=r-_hSanSjLZ_YXFpZUb0Up94u5F8JHp70Vf2nlONPSg,11269
+ai_edge_torch/generative/utilities/autoencoder_loader.py,sha256=G2Nosy33JzkjGALPR4JjvffdFX1JWOj2zjbbuaDJEgg,10065
+ai_edge_torch/generative/utilities/loader.py,sha256=Hs92478j1g4jQGvbdP1aWvOy907HjwqQZE-NFy6HELo,11326
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=guDTv-12UUvJGl4eDvvZX3t4rRKewfXO8SpcYXM6gbc,16156
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
@@ -102,8 +107,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=ExThdTXqnWmGC3-F6sdXbXr8nYzkEe_qCz
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/METADATA,sha256=2yFjQFvs93hoppwWWSJi-B9HC2n-h1s2rzjsyGXU1zI,1748
-ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240606.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240608.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240608.dist-info/METADATA,sha256=eM8KLGmQ4Kc6bVDcSDSJ12O-jFwr0ARV7uGeh-T8nvk,1748
+ai_edge_torch_nightly-0.2.0.dev20240608.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240608.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240608.dist-info/RECORD,,