# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
#
# 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.

"""Dynamo exporter.

Wraps `torch.export.export(strict=False)` with helpers that make Transformers
models exportable. The export pipeline uses five sections, in execution order:

1. **Model signature patch** (`patch_forward_signature`): replaces `model.forward`
   with a flat explicit signature derived from `sample_inputs` so `torch.export` does
   not expand `**kwargs` into a `combined_args` bundle that mismatches `dynamic_shapes`.
   This is the entry contract `torch.export` reads before tracing.
2. **Model patches** (`_PATCHES["dynamo"]` via `apply_patches("dynamo")`): reversible
   class-attribute swaps applied during tracing to replace non-exportable model patterns
   (data-dependent loops, in-place ops, mask checks) with export-safe equivalents.
   Modeling code itself is not updated because these patches are too model-specific.
3. **Pytree registration** (`register_cache_pytrees_for_model`): flatten/unflatten
   hooks (via `torch.utils._pytree.register_pytree_node`) for Cache subclasses and
   custom containers so `torch.export` can trace through them.
4. **Dynamic shapes** (`get_auto_dynamic_shapes`): automatic `Dim.AUTO` inference
   for all tensor and cache inputs when `DynamoConfig.dynamic=True`.
5. **Model state cleanup** (`reset_model_state`): non-Cache stateful module attributes
   (`_STATEFUL_CACHE_ATTRS`) are saved on entry, set to `None` during the trace, and
   restored on exit — so a previous eager forward doesn't leak into the trace and any
   FakeTensors the trace planted are discarded before the next eager forward.
"""

from __future__ import annotations

import copy
import importlib
import inspect
import sys
import types
from collections.abc import MutableMapping
from contextlib import contextmanager
from typing import Any

from ..utils import logging
from ..utils.import_utils import is_detectron2_available, is_torch_available, torch_compilable_check
from .base import HfExporter
from .configs import DynamoConfig
from .utils import apply_patches, patch_attributes, prepare_for_export, register_patch


if is_torch_available():
    import torch
    from torch.export import ExportedProgram

    from ..cache_utils import Cache
    from ..modeling_utils import PreTrainedModel


logger = logging.get_logger(__file__)


class DynamoExporter(HfExporter):
    """Exporter that converts a [`PreTrainedModel`] to an `ExportedProgram`.

    Example:

    ```python
    >>> from transformers.exporters.exporter_dynamo import DynamoExporter, DynamoConfig

    >>> exporter = DynamoExporter()
    >>> exported = exporter.export(model, inputs, config=DynamoConfig(dynamic=True))
    >>> outputs = exported.module()(**inputs)
    ```
    """

    required_packages = ["torch"]
    min_versions = {"torch": "2.11.0"}
    tested_versions = {"torch": "2.12.0"}

    def export(
        self,
        model: PreTrainedModel,
        sample_inputs: MutableMapping[str, Any],
        config: DynamoConfig | dict[str, Any],
    ) -> ExportedProgram:
        if isinstance(config, dict):
            config = DynamoConfig(**config)
        elif not isinstance(config, DynamoConfig):
            raise TypeError(f"Expected config to be a DynamoConfig or dict, got {type(config)}")

        model, sample_inputs, output_flags = prepare_for_export(model, sample_inputs)

        dynamic_shapes = config.dynamic_shapes
        if config.dynamic and dynamic_shapes is None:
            logger.warning_once(
                "`dynamic=True` with no explicit `dynamic_shapes` marks every input axis `Dim.AUTO`, so "
                "torch.export resolves symbolic shapes for all of them — including axes that are actually "
                "fixed (batch, a size-1 decode step, num_heads/head_dim). Passing explicit `dynamic_shapes` "
                "that mark only the axes which vary bypasses that symbolic-shape resolution and exports "
                "significantly faster."
            )
            dynamic_shapes = get_auto_dynamic_shapes(sample_inputs)

        register_cache_pytrees_for_model(model)

        with (
            apply_patches("dynamo"),
            reset_model_state(model),
            patch_model_config(model, output_flags),
            patch_forward_signature(model, sample_inputs),
        ):
            exported_program: ExportedProgram = torch.export.export(
                model,
                args=(),
                kwargs=copy.deepcopy(dict(sample_inputs)),
                strict=config.strict,
                dynamic_shapes=dynamic_shapes,
                prefer_deferred_runtime_asserts_over_guards=config.prefer_deferred_runtime_asserts_over_guards,
            )

        return exported_program


# ── Stage 1: Model signature patch ──────────────────────────────────────────
# Replaces `model.forward` with a flat explicit signature derived from the
# inputs dict so `torch.export` does not expand `**kwargs` into a large bundle.
# `patch_model_config` lives here too — it strips output flags from the inputs
# and applies them onto `model.config` for the duration of the trace.


# Output flags stripped from inputs and applied onto `model.config` for the trace.
@contextmanager
def patch_model_config(model: PreTrainedModel, output_flags: dict[str, Any]):
    """Reversibly tweak `model.config` for the trace:

    - Applies `output_flags` (popped from inputs by `prepare_for_export`) onto
      `model.config.<flag>` so the model picks them up via its usual `<flag> if <flag> is
      not None else self.config.<flag>` fallback.
    - Disables `use_mamba_kernels` on every submodel's config that declares it (mamba/jamba
      kernels are not exportable).

    Originals are restored on exit. Flags whose value is `None`, or that the config doesn't
    declare, are silently skipped — useful for submodels that don't accept every parent flag.
    """
    config_patches = []
    for flag, value in output_flags.items():
        if value is None or not hasattr(model, "config") or not hasattr(model.config, flag):
            continue
        config_patches.append((model.config, flag, lambda _original, v=value: v))
    for module in model.modules():
        if hasattr(module, "config") and hasattr(module.config, "use_mamba_kernels"):
            config_patches.append((module.config, "use_mamba_kernels", lambda _original: False))
    with patch_attributes(config_patches):
        yield


@contextmanager
def patch_forward_signature(model: PreTrainedModel, inputs: dict[str, Any]):
    """Temporarily replace `model.forward` with a flat explicit signature derived from `inputs`.

    `torch.export` infers the exported function signature from `model.forward.__signature__`.
    Most transformers models use `**kwargs: Unpack[TransformersKwargs]`, which causes
    `torch.export` to expand the signature into a large `combined_args` bundle that
    mismatches the `dynamic_shapes` dict. This patch replaces the forward with a
    minimal signature containing only the keys present in `inputs`.
    """
    original_forward = model.forward

    def _flat_forward(**kwargs):
        return original_forward(**kwargs)

    _flat_forward.__signature__ = inspect.Signature(
        [inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None) for k in inputs]
    )

    try:
        model.forward = _flat_forward
        yield
    finally:
        model.forward = original_forward


# ── Stage 2: Model patches ────────────────────────────────────────────────────
# Reversible class-attribute swaps applied during `torch.export` tracing via
# `apply_patches("dynamo")`. Each replaces a non-exportable model pattern
# (data-dependent control flow, in-place ops on views, etc.) with an
# export-safe equivalent on the owning class — every live instance sees the
# replacement until the context exits. Modeling code itself is not updated
# because these patches are too model-specific; we do strive to keep modeling
# code compliant where reasonable.
#
# Each `@register_patch("dynamo", *dotted_paths)` decorator targets one or
# more `Class.method` paths and wraps a `factory(original) -> replacement`.
# Multiple paths share the same factory when the same method shape needs to be
# swapped across several classes (e.g. `_reshaped_vision_attention_forward`
# applied to every chunked-vision attention class — see the long list below).


@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier")
def _patch_classifier_cast(_original):
    """Disable classifier dtype cast in nllb-moe (not traceable)."""
    return lambda self, *args, **kwargs: None


@register_patch("dynamo", "torch.nn.functional.scaled_dot_product_attention")
def _patch_sdpa(original):
    """Route SDPA through the MATH backend on CPU during tracing — CPU SDPA's flash/efficient
    paths guard on ``Eq(batch, 1)`` (upstream https://github.com/pytorch/pytorch/issues/180202),
    which trips ``GuardOnDataDependentSymNode`` whenever the batch dim comes from a data-dependent
    op like ``pixel_values[bool_mask]`` (Idefics2/3 and most VLMs). The MATH decomposition has no
    batch-1 dispatch, so the guard never fires. CUDA exports are left alone — the GPU kernels
    don't have this guard, and we want the flash/efficient decompositions there.
    """
    from torch.nn.attention import SDPBackend, sdpa_kernel

    def patch(query, *args, **kwargs):
        if query.device.type == "cpu":
            with sdpa_kernel(SDPBackend.MATH):
                return original(query, *args, **kwargs)
        return original(query, *args, **kwargs)

    return patch


@register_patch(
    "dynamo",
    # Canonical definition + the public re-export.
    "transformers.utils.import_utils.is_kernels_available",
    "transformers.utils.is_kernels_available",
    # Local `from ...utils import is_kernels_available` rebinds in modeling modules
    # — each one needs its own override since the name is looked up there.
    "transformers.modeling_utils.is_kernels_available",
    "transformers.models.sam3_video.modeling_sam3_video.is_kernels_available",
    "transformers.models.mra.modeling_mra.is_kernels_available",
    "transformers.models.rwkv.modeling_rwkv.is_kernels_available",
    "transformers.models.yoso.modeling_yoso.is_kernels_available",
)
def _patch_is_kernels_available(_original):
    """Force-disable the optional ``kernels`` library during export — its kernels
    call into native code that ``torch.export`` cannot trace, and the pure-PyTorch
    fallbacks in each model are always traceable."""
    return lambda *args, **kwargs: False


# --- Chunked vision/audio attention ─────────────────────────────────────────
# Sub-encoders that pack multiple variable-length sequences into one flat tensor
# with `cu_seqlens` markers fall back to `split → per-segment SDPA → cat` in the
# unpatched forward, which is a Python loop that `torch.export` can't trace.
# `_reshaped_vision_attention_forward` replaces that loop with a reshape into a
# per-segment batch followed by a single SDPA call. It handles the layout
# differences across encoders (combined `qkv` vs separate `q/k/v` vs separate
# `q_proj/k_proj/v_proj`, asymmetric `q_dim/kv_dim` split, `(cos, sin)` vs single
# rotary tensor vs none, `.proj` vs `.out_proj`, NaViT `(1, T, D)` packing,
# tuple vs single return). The `returns_tuple` flag is bound once per class at
# install time by inspecting the original `forward`'s source.
#
# NOTE: this whole stack of patches becomes unnecessary once transformers adopts a
# proper varlen-attention op (e.g. PyTorch's `torch._nested.scaled_dot_product_attention`
# or a Flex-Attention varlen kernel) — the modeling forwards can then express the
# segmented attention directly with `cu_seqlens` and trace through `torch.export`
# without this reshape-into-batch workaround. Drop this section when that lands.


def _reshaped_vision_attention_forward(
    self,
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor,
    rotary_pos_emb: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
    returns_tuple: bool = False,
    **kwargs,
):
    """Export-safe chunked vision/audio attention: reshape segments into a batch dim,
    apply rotary if provided, run one SDPA call, project, and re-emit in the original layout."""

    # Normalise NaViT-style `(1, T, D)` packing (minicpmv4_6) to the flat `(T, D)` layout
    # the rest of this wrapper assumes. The leading dim is always 1 — multi-image batches
    # are packed along the sequence dim.
    needs_batch_restore = hidden_states.ndim == 3
    if needs_batch_restore:
        hidden_states = hidden_states.squeeze(0)

    seq_length = hidden_states.shape[0]
    torch_compilable_check(
        seq_length != 0,
        "Chunked vision attention received an empty input.",
    )
    num_segments = cu_seqlens.shape[0] - 1
    # The reshape-into-batch below needs equal-length segments.
    segment_lengths = cu_seqlens[1:] - cu_seqlens[:-1]
    torch_compilable_check(
        (segment_lengths == segment_lengths[0]).all(),
        "Chunked vision attention requires uniform segment lengths during export. "
        "Ensure all images have the same resolution (use do_resize=True in the processor) "
        "or pad inputs to a common size.",
    )

    if hasattr(self, "qkv"):
        # Grouped-query attention (q_dim != kv_dim, e.g. Exaone4.5) splits asymmetrically;
        # uniform reshape into (seq, 3, num_heads, -1) only works when Q, K, V share the head count.
        if hasattr(self, "q_dim") and hasattr(self, "kv_dim") and self.q_dim != self.kv_dim:
            query_states, key_states, value_states = self.qkv(hidden_states).split(
                [self.q_dim, self.kv_dim, self.kv_dim], dim=-1
            )
            query_states = query_states.view(seq_length, self.num_heads, self.head_dim)
            key_states = key_states.view(seq_length, self.num_key_value_heads, self.head_dim)
            value_states = value_states.view(seq_length, self.num_key_value_heads, self.head_dim)
        else:
            query_states, key_states, value_states = (
                self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).transpose(0, 1).unbind(0)
            )
    else:
        q_proj = getattr(self, "q_proj", getattr(self, "q", None))
        k_proj = getattr(self, "k_proj", getattr(self, "k", None))
        v_proj = getattr(self, "v_proj", getattr(self, "v", None))
        query_states = q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)
        key_states = k_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)
        value_states = v_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim)

    if position_embeddings is not None:
        # Each vision encoder ships its own ``apply_rotary_pos_emb_vision`` in its modeling file
        # (Qwen2-VL's takes (q, k, cos, sin), Qwen2.5/3-Omni's takes (x, rotary_emb), etc.). Look
        # it up on the model's own module so this patch stays signature-agnostic across the
        # ~19 attention classes it's installed on.
        apply_rotary_pos_emb_vision = sys.modules[type(self).__module__].apply_rotary_pos_emb_vision
        if isinstance(position_embeddings, (tuple, list)):
            # (cos, sin) tuple convention — most VL encoders.
            cos, sin = position_embeddings
            query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
        else:
            # Single `rotary_pos_emb` tensor convention — Qwen2.5/3 Omni vision applies rotary per-states.
            query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), position_embeddings).squeeze(0)
            key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), position_embeddings).squeeze(0)

    seg_len = seq_length // num_segments

    # (seq, heads, dim) → (n_seg, seg_len, heads, dim) → (n_seg, heads, seg_len, dim)
    def _to_batched(t):
        return t.unflatten(0, (num_segments, seg_len)).transpose(1, 2)

    query_states = _to_batched(query_states)
    key_states = _to_batched(key_states)
    value_states = _to_batched(value_states)

    torch_compilable_check(query_states.shape[0] != 0, "Reshaped chunked-vision attention got zero batch.")
    torch_compilable_check(query_states.shape[2] != 0, "Reshaped chunked-vision attention got zero seq.")
    attn_output = torch.nn.functional.scaled_dot_product_attention(
        query_states,
        key_states,
        value_states,
        is_causal=False,
        scale=self.scaling,
        dropout_p=0.0 if not self.training else self.attention_dropout,
        enable_gqa=getattr(self, "num_key_value_heads", self.num_heads) != self.num_heads,
    )

    # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim).
    attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1)
    out_proj = self.proj if hasattr(self, "proj") else self.out_proj
    attn_output = out_proj(attn_output)

    if needs_batch_restore:
        attn_output = attn_output.unsqueeze(0)

    return (attn_output, None) if returns_tuple else attn_output


@register_patch(
    "dynamo",
    # Combined `qkv` + `(cos, sin)` rotary + `.proj`
    "transformers.models.qwen2_vl.modeling_qwen2_vl.VisionAttention.forward",
    "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLVisionAttention.forward",
    "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward",
    "transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe.Qwen3VLMoeVisionAttention.forward",
    "transformers.models.qwen3_5.modeling_qwen3_5.Qwen3_5VisionAttention.forward",
    "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe.Qwen3_5MoeVisionAttention.forward",
    "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeVisionAttention.forward",
    "transformers.models.glm4v.modeling_glm4v.Glm4vVisionAttention.forward",
    "transformers.models.glm4v_moe.modeling_glm4v_moe.Glm4vMoeVisionAttention.forward",
    "transformers.models.glm_ocr.modeling_glm_ocr.GlmOcrVisionAttention.forward",
    "transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe.Ernie4_5_VLMoeVisionAttention.forward",
    # Combined `qkv` + optional `(cos, sin)` rotary + `.proj`
    "transformers.models.cohere_compass.modeling_cohere_compass.CohereCompassVisionAttention.forward",
    # Asymmetric `qkv` split + `(cos, sin)` rotary + `.proj`
    "transformers.models.exaone4_5.modeling_exaone4_5.Exaone4_5_VisionAttention.forward",
    # Separate `.q` / `.k` / `.v` + single rotary tensor + `.proj`
    "transformers.models.qwen2_5_omni.modeling_qwen2_5_omni.Qwen2_5OmniVisionAttention.forward",
    # Separate `q_proj`/`k_proj`/`v_proj` + `(cos, sin)` rotary + `.proj` (single return)
    "transformers.models.kimi_k25.modeling_kimi_k25.Kimi_K25VisionAttention.forward",
    "transformers.models.muse_glimmer.modeling_muse_glimmer.MuseGlimmerVisionAttention.forward",
    # Separate `_proj` + `(cos, sin)` rotary + `.out_proj` (tuple return)
    "transformers.models.video_llama_3.modeling_video_llama_3.VideoLlama3VisionAttention.forward",
    "transformers.models.paddleocr_vl.modeling_paddleocr_vl.PaddleOCRVisionAttention.forward",
    # NaViT (1, T, D) + separate `_proj` + `.out_proj` (tuple return)
    "transformers.models.minicpmv4_6.modeling_minicpmv4_6.MiniCPMV4_6VisionAttention.forward",
    # Audio attention: separate `_proj` + `.out_proj`, no rotary
    "transformers.models.qwen2_5_omni.modeling_qwen2_5_omni.Qwen2_5OmniAudioAttention.forward",
    "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeAudioAttention.forward",
    "transformers.models.qwen3_asr.modeling_qwen3_asr.Qwen3ASRAudioAttention.forward",
)
def _patch_chunked_vision_attention(original):
    """Bind `returns_tuple` once per class by inspecting the original forward's source."""
    src = inspect.getsource(original)
    returns_tuple = "return attn_output, attn_weight" in src or "return attn_output, None" in src

    def forward(self, *args, **kwargs):
        return _reshaped_vision_attention_forward(self, *args, returns_tuple=returns_tuple, **kwargs)

    return forward


# ── Stage 3: Pytree registration ─────────────────────────────────────────────
# torch.export needs pytree flatten/unflatten for Cache objects and other
# custom types. The generic flattener serialises any object to a JSON-native
# context (bools, ints, strings, dicts, lists) while collecting tensors into
# a flat list — the inverse reconstructs the original object.
#
# To register a new type: it should be handled automatically by the generic
# flattener. If not, add a branch in _flatten_to_context / _unflatten_from_context.


def _class_to_path(cls: type) -> str:
    return f"{cls.__module__}:{cls.__qualname__}"


def _path_to_class(path: str) -> type:
    module_name, qualname = path.split(":", 1)
    obj = importlib.import_module(module_name)
    for part in qualname.split("."):
        obj = getattr(obj, part)
    return obj


def _maybe_sym_constant(sym: Any) -> Any:
    """The concrete value a ``Sym*`` has already specialized to, or ``None`` if it's still dynamic.
    Each ``Sym*`` type has its own accessor (``maybe_as_bool``/``maybe_as_float``/``maybe_as_int``);
    there is no generic one."""
    node = sym.node
    if isinstance(sym, torch.SymBool):
        return node.maybe_as_bool()
    if isinstance(sym, torch.SymFloat):
        return node.maybe_as_float()
    return node.maybe_as_int()


def _flatten_to_context(obj: Any, tensors: list) -> Any:
    """Single-pass: recursively build a JSON-native context while collecting tensors into `tensors`."""
    # --- Pure Python / JSON-native (exact type check — subclasses fall through to stateful objects) ---
    if obj is None or type(obj) in (bool, int, float, str):
        return obj
    if type(obj) is list:
        return [_flatten_to_context(i, tensors) for i in obj]
    if type(obj) is dict:
        return {k: _flatten_to_context(v, tensors) for k, v in obj.items()}

    # --- Torch objects ---
    if isinstance(obj, torch.Tensor):
        idx = len(tensors)
        tensors.append(obj)
        return {"_t": "tensor", "i": idx}
    if isinstance(obj, torch.Size):
        return {"_t": "size", "v": list(obj)}
    if isinstance(obj, torch.device):
        return {"_t": "device", "s": str(obj)}
    if isinstance(obj, torch.dtype):
        return {"_t": "dtype", "n": str(obj).removeprefix("torch.")}
    if isinstance(obj, torch.layout):
        return {"_t": "layout", "n": str(obj).removeprefix("torch.")}
    if isinstance(obj, (torch.SymInt, torch.SymFloat, torch.SymBool)):
        # A Sym* that has already specialized to a concrete constant is not a genuine dynamic
        # graph output — bake it as a plain scalar instead of a leaf. Leaving it as a leaf makes
        # flatten non-deterministic: the same field can be a constant-valued SymInt at one trace
        # point (the dynamo out_spec capture) and an already-materialized python scalar at another
        # (the aot-decomposition retrace), so the leaf count flips and `treespec.unflatten` fails
        # with an off-by-one. deepseek_v4 hits this via two sibling cache layers
        # (DeepseekV4HCACache / DeepseekV4CSACache) sharing a `cumulative_length` counter that one
        # path leaves as a constant SymInt and the other as a python int.
        const = _maybe_sym_constant(obj)
        if const is not None:
            return const
        idx = len(tensors)
        tensors.append(obj)
        return {"_t": "sym", "i": idx}

    # --- Python types ---
    if isinstance(obj, type):
        return {"_t": "type", "p": _class_to_path(obj)}

    # --- Generic Python objects (by structural category) ---
    cls = type(obj)
    if isinstance(obj, dict):  # dict subclasses (OrderedDict, etc.)
        return {
            "_t": "map",
            "p": _class_to_path(cls),
            "v": {k: _flatten_to_context(v, tensors) for k, v in obj.items()},
        }
    if isinstance(obj, (tuple, list, set, frozenset)):  # sequences/sets incl. NamedTuple
        return {
            "_t": "seq",
            "p": _class_to_path(cls),
            "v": [_flatten_to_context(i, tensors) for i in obj],
        }
    if isinstance(obj, types.MethodType):
        # A bound method can't be flattened into pytree context. Models shouldn't bind methods onto
        # objects they return at forward time (e.g. recurrent_gemma binds `get_seq_length`/
        # `get_mask_sizes` onto its `DynamicCache`) — that pattern isn't exportable; such a model is
        # skipped until the binding is refactored away (e.g. into a `Cache` subclass).
        raise TypeError("Cannot flatten a bound method for pytree context")
    if hasattr(obj, "__dict__"):
        state = {k: _flatten_to_context(v, tensors) for k, v in vars(obj).items()}
        return {"_t": "obj", "p": _class_to_path(cls), "s": state}

    raise TypeError(f"Cannot flatten {type(obj).__name__} for pytree context")


def _unflatten_from_context(ctx: Any, tensors: list) -> Any:
    """Reconstruct an object from its JSON-native context, substituting tensor index markers."""
    # --- Pure Python / JSON-native ---
    if ctx is None or type(ctx) in (bool, int, float, str):
        return ctx
    if type(ctx) is list:
        return [_unflatten_from_context(i, tensors) for i in ctx]
    if type(ctx) is dict and "_t" not in ctx:
        return {k: _unflatten_from_context(v, tensors) for k, v in ctx.items()}

    # --- Torch objects ---
    t = ctx["_t"]
    if t == "tensor":
        return tensors[ctx["i"]]
    if t == "layout":
        return getattr(torch, ctx["n"])
    if t == "dtype":
        return getattr(torch, ctx["n"])
    if t == "device":
        return torch.device(ctx["s"])
    if t == "size":
        return torch.Size(ctx["v"])
    if t == "sym":
        return tensors[ctx["i"]]

    # --- Python types ---
    if t == "type":
        return _path_to_class(ctx["p"])

    # --- Generic Python objects ---
    if t == "map":
        cls = _path_to_class(ctx["p"])
        return cls({k: _unflatten_from_context(v, tensors) for k, v in ctx["v"].items()})
    if t == "seq":
        cls = _path_to_class(ctx["p"])
        items = [_unflatten_from_context(i, tensors) for i in ctx["v"]]
        try:
            return cls(items)  # tuple, list subclass, set, frozenset, etc.
        except TypeError:
            return cls(*items)  # NamedTuple (requires positional args)
    if t == "obj":
        cls = _path_to_class(ctx["p"])
        instance = cls.__new__(cls)
        for k, v in ctx["s"].items():
            instance.__dict__[k] = _unflatten_from_context(v, tensors)
        return instance

    raise TypeError(f"Unknown tag {t!r} in pytree context")


def _pytree_flatten(obj: Any) -> tuple[list, Any]:
    tensors: list = []
    context = _flatten_to_context(obj, tensors)
    return tensors, context


def _pytree_flatten_with_keys(obj: Any):
    leaves, context = _pytree_flatten(obj)
    return [(torch.utils._pytree.SequenceKey(i), leaf) for i, leaf in enumerate(leaves)], context


def _pytree_unflatten(values, context: Any) -> Any:
    return _unflatten_from_context(context, list(values))


def register_pytree_node(object_cls: type):
    """Register a single class (e.g. a `Cache` subclass like `StaticCache`) as a torch.export pytree
    node, so `torch.export.load` can unflatten it as a graph input without needing the original model."""
    try:
        torch.utils._pytree.register_pytree_node(
            object_cls,
            _pytree_flatten,
            _pytree_unflatten,
            serialized_type_name=_class_to_path(object_cls),
            flatten_with_keys_fn=_pytree_flatten_with_keys,
        )
    except ValueError as e:
        if "already registered as pytree node" not in str(e):
            raise


def _iter_subclasses(cls: type):
    for subclass in cls.__subclasses__():
        yield subclass
        yield from _iter_subclasses(subclass)


def register_cache_pytrees_for_model(model: PreTrainedModel):
    """Register all relevant cache types as pytree nodes for torch.export."""
    # All transformers Cache subclasses
    for cache_type in _iter_subclasses(Cache):
        register_pytree_node(cache_type)

    # Model-specific cache classes not inheriting from Cache (e.g. custom per-model caches)
    for _, obj in inspect.getmembers(inspect.getmodule(model)):
        if (
            inspect.isclass(obj)
            and obj.__module__ == model.__class__.__module__
            and obj.__name__.endswith("Cache")
            and not issubclass(obj, Cache)
        ):
            register_pytree_node(obj)

    # detectron2 ImageList (used by layoutlmv2)
    if is_detectron2_available() and isinstance(model, PreTrainedModel) and model.config.model_type == "layoutlmv2":
        from detectron2.structures.image_list import ImageList

        register_pytree_node(ImageList)


# ── Stage 4: Dynamic shapes ─────────────────────────────────────────────────
# Automatic `Dim.AUTO` inference for all tensor and cache inputs when
# `DynamoConfig.dynamic` is True and no explicit `dynamic_shapes` are provided.


def _auto_dynamic_shape(tensor: torch.Tensor) -> dict[int, torch.export.Dim]:
    """Generate a dynamic shape with all dimensions set to Dim.AUTO for a given tensor."""
    return dict.fromkeys(range(tensor.dim()), torch.export.Dim.AUTO)


def get_auto_dynamic_shapes(inputs: Any) -> Any:
    """Recursively build dynamic shapes for any input value.

    - Tensors → per-dimension Dim.AUTO spec.
    - Scalars / None → None (no dynamic dims).
    - Objects with ``__dict__`` (ModelOutput, Cache, …) → flat list of leaf specs,
      matching the ``TreeSpec(list, …)`` that torch.export produces for these types.
    - Lists / tuples → same container type, recursed element-wise.
    - Plain dicts → recursed dict of specs.
    - Everything else → None.
    """
    if isinstance(inputs, torch.Tensor):
        return _auto_dynamic_shape(inputs)
    if inputs is None or isinstance(inputs, (int, float, bool, str)):
        return None
    if hasattr(inputs, "__dict__"):
        leaves, _ = _pytree_flatten(inputs)
        return get_auto_dynamic_shapes(leaves)
    if type(inputs) in (list, tuple, set, frozenset):
        return type(inputs)(get_auto_dynamic_shapes(v) for v in inputs)
    if type(inputs) is dict:
        return {k: get_auto_dynamic_shapes(v) for k, v in inputs.items()}
    return None


# ── Stage 5: Model state cleanup ────────────────────────────────────────────
# `torch.export` traces forward with FakeTensors, which can leave non-Cache stateful
# tensor attributes as FakeTensors after tracing — a follow-up eager forward then
# hits shape/dtype mismatches when it reuses the stale state. We also want stale
# eager-mode state cleared on entry so it doesn't leak into the trace.
# `reset_model_state` brackets the `torch.export.export` call: it saves every
# attribute in `_STATEFUL_CACHE_ATTRS` on every submodule, sets them to `None` for
# the trace, and restores the originals on exit (finally semantics).
#
# To register a new stateful attribute: append its name to `_STATEFUL_CACHE_ATTRS`.

_STATEFUL_CACHE_ATTRS = (
    "cached_rotary_positional_embedding",  # wav2vec2_bert, seamless_m4t, clvp
    "cached_sequence_length",  # wav2vec2_bert, seamless_m4t, clvp
)


@contextmanager
def reset_model_state(model: torch.nn.Module):
    """Save each `_STATEFUL_CACHE_ATTRS` value, null it for the trace, restore on exit.

    FakeTensors that `torch.export` plants into these attributes during the trace are
    discarded by the restore.
    """
    originals = [
        (module, attr, getattr(module, attr))
        for module in model.modules()
        for attr in _STATEFUL_CACHE_ATTRS
        if hasattr(module, attr)
    ]
    for module, attr, _ in originals:
        setattr(module, attr, None)
    try:
        yield
    finally:
        for module, attr, original in originals:
            setattr(module, attr, original)