Moved types.py to generative/utilities.

PiperOrigin-RevId: 746151614

Author: haozha111

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -20,7 +20,7 @@
 
 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

ai_edge_torch/generative/layers/kv_cache.py

@@ -20,7 +20,7 @@
 
 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

@@ -19,8 +19,8 @@
 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
-from ai_edge_torch.generative.layers.experimental import types
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import types
 from multipledispatch import dispatch
 import torch