in optimum/exporters/executorch/integrations.py [0:0]
def _prepare_export_inputs(self):
"""
Prepare example inputs and configurations for export.
Returns:
example_input_ids (torch.Tensor): Example input IDs tensor.
example_cache_position (torch.Tensor): Example cache position tensor.
dynamic_shapes (dict or None): Dynamic shape specifications for export.
strict (bool): Whether to use strict export mode.
"""
# Default values for legacy or fallback cases
example_input_ids = torch.tensor([[1]], dtype=torch.long)
example_cache_position = torch.tensor([0], dtype=torch.long)
dynamic_shapes = None
strict = True
is_using_hybrid_cache_wo_custom_sdpa_kv_cache = (
hasattr(self.config, "layer_types")
and getattr(self.config, "sliding_window", None) is not None
and not (self.use_custom_kv_cache and self.use_custom_sdpa)
)
if is_transformers_version(">", "4.52.0") and not is_using_hybrid_cache_wo_custom_sdpa_kv_cache:
# Prepare inputs with dynamic shapes
seq_length = 3 # Sequence length > 1 to avoid specialization issues
example_input_ids = torch.zeros((1, seq_length), dtype=torch.long)
example_cache_position = torch.arange(seq_length, dtype=torch.long)
max_seq_len = self.metadata.get("get_max_seq_len")
sliding_window = self.metadata.get("sliding_window", float("inf"))
max_dim = min(max_seq_len, sliding_window) - 1
seq_len_dim = torch.export.Dim("seq_length_dim", max=max_dim)
dynamic_shapes = {
"input_ids": {1: seq_len_dim},
"cache_position": {0: seq_len_dim},
}
strict = parse(torch.__version__) != parse("2.7.0") # Workaround for PyTorch bug #150994
return example_input_ids, example_cache_position, dynamic_shapes, strict