hyperonnx 1.0.2__tar.gz → 1.0.3__tar.gz

hyperonnx-1.0.3/.claude/CLAUDE.md

@@ -0,0 +1,27 @@
+# HyperONNX Instructions
+
+## Workflow
+
+- Check relevant skills in `.claude/skills/` before handling specialized tasks.
+- Prefer minimal, targeted changes that preserve existing project structure and style.
+- Validate the files you change with the smallest relevant command before finishing.
+
+## Python Changes
+
+- Follow the existing Python style already used in `hyperonnx/`.
+- When Python files are modified, use the lint skill workflow and prefer `uv run ruff check hyperonnx`.
+- Format Python changes with `uv run ruff format hyperonnx` when formatting is needed.
+
+## Version Compatibility (Hard Constraints)
+
+⚠️ **Supported Versions**: This project is currently tested and working on:
+- **torch < 2.11** (torch 2.5–2.10 verified)
+- **onnxscript < 0.6.0**
+
+Both torch >= 2.11 and onnxscript >= 0.6.0 introduce breaking changes to module structure, dtype handling, and ONNX IR that are not yet resolved. **Do not use these newer versions** without explicit future work to support them.
+
+**When users ask about newer versions**: Acknowledge the constraint, suggest using compatible versions from `pyproject.toml` optional dependencies, and note that torch >= 2.11 support is a known gap requiring future implementation.
+
+## Commits
+
+- If asked to create a commit, use a clear and concise commit message that describes the actual change.

hyperonnx-1.0.3/.claude/skills/torch-dynamo-onnx-export/SKILL.md

@@ -0,0 +1,422 @@
+---
+name: torch-dynamo-onnx-export
+description: Use this skill when exporting PyTorch models to ONNX using the modern torch.export (dynamo) path in HyperONNX. Best for models with dynamic control flow, complex tensor operations, and hierarchical custom operators defined via onnxscript.
+argument-hint: Describe the model, custom operations (if any), target opset, and whether you need hierarchical export or simple custom operator definitions.
+---
+
+# Torch Dynamo ONNX Export
+
+Use this skill when the task is about the modern Torch ONNX export path in HyperONNX, meaning the `torch.export`-based path with `dynamo=True`.
+
+This skill is specific to this repository. It captures the actual workflow used by HyperONNX code and tests, rather than generic PyTorch ONNX advice.
+
+## Torch Behaviour
+
+- PyTorch 2.6+ introduced `torch.export` for AOT (Ahead-of-Time) tracing, enabling more dynamic models than TorchScript.
+- PyTorch 2.9+ made `dynamo=True` the default in `torch.onnx.export()`.
+- The dynamo path is based on FX graphs and bytecode analysis, not static tracing.
+- On PyTorch 2.5 and below, `dynamo` parameter may not exist; always check `torch.onnx.export` signature and use the wrapper `torch_export_handle_lower_version` from HyperONNX.
+- In this repository, prefer being explicit: always pass `dynamo=True` to clarify intent when using the modern exporter.
+
+⚠️ **Version Constraint**: This skill is tested and supported on **torch < 2.11** with **onnxscript < 0.6.0**. These versions have breaking changes in ONNX/module structure that require future work to resolve.
+
+## What This Skill Produces
+
+Produce one of these outcomes:
+
+- A correct export snippet using the modern `torch.export` path with `dynamo=True`.
+- A HyperONNX code change that leverages `torch.export` for better dynamic model support.
+- A custom operator definition using onnxscript and custom_translation_table for the dynamo path.
+- A debugging checklist for why a dynamo export is failing or behaving unexpectedly.
+- A migration guide from legacy (dynamo=False) to modern (dynamo=True) export.
+
+## When To Use It
+
+Use this skill when the user asks for any of the following:
+
+- Use `torch.onnx.export` with the modern or dynamo path.
+- Force torch.export instead of TorchScript tracing.
+- Export a model with `dynamo=True`.
+- Create custom ONNX operators using onnxscript for the dynamo path.
+- Build a custom_translation_table for hierarchical export.
+- Understand why dynamo export behaves differently from legacy export.
+- Export models with dynamic shapes, loops, or data-dependent control flow.
+
+Do not use this skill for generic ONNX export or torch.export advice unrelated to HyperONNX.
+
+## Repository Facts To Ground Decisions
+
+- HyperONNX exposes `torch_export_handle_lower_version` in [hyperonnx/torch_export.py](../../../hyperonnx/torch_export.py) to bridge version differences.
+- The main hierarchical export entrypoint is [hyperonnx/hyper_export.py](../../../hyperonnx/hyper_export.py).
+- The dynamo-specific branch for hierarchical replacement lives in [hyperonnx/exporter/dynamo.py](../../../hyperonnx/exporter/dynamo.py).
+- Custom operator definition using onnxscript is built dynamically in [hyperonnx/exporter/dynamo.py](../../../hyperonnx/exporter/dynamo.py) via `build_onnxscript` and IRBuilder.
+- Transformer attention for the dynamo path is handled in [hyperonnx/transformers/attention.py](../../../hyperonnx/transformers/attention.py) via `attention_translation_table()`.
+- Reference tests: [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py), [tests/expoter/test_dynamo_build_onnxscript.py](../../../tests/expoter/test_dynamo_build_onnxscript.py), [tests/expoter/test_dynamo_make_custom_op.py](../../../tests/expoter/test_dynamo_make_custom_op.py).
+
+## Dynamo Export Fundamentals
+
+### Core Concepts
+
+**torch.export**: Produces an ExportedProgram (AOT traced graph) that preserves more semantics than TorchScript, including data-dependent control flow and dynamic shapes.
+
+**custom_translation_table**: A dictionary mapping PyTorch callables (e.g., custom ops, modules) to their ONNX counterparts (via onnxscript). The dynamo exporter uses this table to translate FX nodes to ONNX operators.
+
+**onnxscript**: A framework for defining ONNX functions programmatically. HyperONNX uses `onnxscript.OnnxFunction` and `irbuilder.IRFunction` to build custom operators dynamically.
+
+**replace_with_custom_op**: Context manager in [hyperonnx/exporter/dynamo.py](../../../hyperonnx/exporter/dynamo.py) that builds a custom_translation_table from module_spec, yielding it for use in export.
+
+**build_onnxscript**: Dynamically constructs an ONNX function from a module's spec (inputs, outputs, dtype info). Used to create custom operator definitions without writing onnxscript manually.
+
+### Key Differences from Legacy Export (TorchScript)
+
+| Aspect | Legacy (dynamo=False) | Dynamo (dynamo=True) |
+|--------|----------------------|---------------------|
+| **Tracer** | TorchScript (static tracing) | torch.export (FX bytecode) |
+| **Dynamic Shapes** | Limited | Better supported |
+| **Control Flow** | Loops unrolled, if/else lost | Can be preserved (with torch.cond) |
+| **Custom Op Registration** | torch.onnx.register_custom_op_symbolic | custom_translation_table + onnxscript |
+| **Module Replacement** | DuckForward (torch.autograd.Function) | Direct FX node replacement |
+| **Symbolic Gen** | Automatic from DuckForward.symbolic | Explicit via onnxscript |
+| **Attention Registration** | register_attention_opsets() | attention_translation_table() |
+| **Performance** | Often faster export | Slower export (bytecode analysis) |
+
+## Defining Custom Operators for Dynamo Export
+
+### Pattern 1: Direct onnxscript function (simple custom op)
+
+For single custom operations without hierarchical wrapping:
+
+```python
+import onnxscript
+import torch
+
+# Step 1: Define an onnxscript function
+_MY_OPSET = onnxscript.values.Opset("custom_domain", 1)
+
+@onnxscript.script(_MY_OPSET)
+def my_custom_op(x: onnxscript.FLOAT) -> onnxscript.FLOAT:
+    return x * 2.0 + 1.0
+
+# Step 2: Build custom_translation_table
+custom_translation_table = {
+    torch.ops.custom_lib.my_op: my_custom_op,
+}
+
+# Step 3: Export with dynamo=True and the table
+torch.onnx.export(
+    model,
+    example_inputs,
+    'model.onnx',
+    dynamo=True,
+    custom_translation_table=custom_translation_table,
+    opset_version=18,
+)
+```
+
+Key rules:
+- The onnxscript function signature must match the PyTorch operation signature.
+- The dictionary key must be a callable (e.g., torch.ops, a custom op).
+- Always pass opset_version to align with onnxscript opset expectations.
+
+### Pattern 2: HyperONNX hierarchical export with dynamic onnxscript
+
+For complex modules that need hierarchical export, use `replace_with_custom_op` and `build_onnxscript`:
+
+```python
+import torch
+from hyperonnx import export_hyper_onnx
+
+class MyCustomModule(torch.nn.Module):
+    """A complex module with loops or dynamic shapes."""
+    def forward(self, x):
+        for i in range(x.shape[0]):
+            x = x * 2
+        return x
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(10, 20),
+    MyCustomModule(),
+    torch.nn.Linear(20, 5),
+)
+
+export_hyper_onnx(
+    model,
+    (torch.randn(1, 10),),
+    'model.onnx',
+    hiera=[MyCustomModule],  # Export as hierarchical function
+    dynamo=True,             # Use torch.export, not TorchScript
+    opset_version=18,
+)
+```
+
+How it works:
+1. HyperONNX traces the model and caches module specs (inputs, outputs, dtypes).
+2. `build_onnxscript` dynamically creates an ONNX function from the module spec.
+3. `replace_with_custom_op` builds the custom_translation_table.
+4. torch.onnx.export (with dynamo=True) translates the FX graph using the table.
+
+Reference: [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py).
+
+### Pattern 3: Attention-specific translation for dynamo
+
+For transformer attention layers:
+
+```python
+import torch
+from hyperonnx.transformers.attention import attention_translation_table
+
+# Step 1: Get the translation table for dynamo path
+translation_table = attention_translation_table()
+
+# Step 2: Export with dynamo=True and the table
+torch.onnx.export(
+    model,
+    example_inputs,
+    'model.onnx',
+    opset_version=24,
+    dynamo=True,
+    custom_translation_table=translation_table,
+)
+```
+
+This uses the modern onnxscript-based Attention definition. See [hyperonnx/transformers/attention.py](../../../hyperonnx/transformers/attention.py) for implementation.
+
+## Building Custom Operators with onnxscript
+
+### Overview
+
+HyperONNX's `build_onnxscript` function constructs ONNX operators automatically from module specs, eliminating the need to write onnxscript manually. However, you can also write custom onnxscript functions for fine-grained control.
+
+### Auto-generation via build_onnxscript
+
+```python
+from hyperonnx.exporter.dynamo import build_onnxscript
+from hyperonnx.typing import default_module_spec, ModuleSpec
+
+# Assume module_spec is populated via trace_module_spec
+spec: ModuleSpec = {...}  # Contains: args, output, signature, etc.
+
+# Auto-build onnxscript function
+onnx_func = build_onnxscript(spec)
+
+# Use in custom_translation_table
+custom_translation_table = {
+    MyModule: onnx_func,
+}
+```
+
+### Manual onnxscript Definition
+
+For more control, write onnxscript directly:
+
+```python
+import onnxscript
+
+# Define a custom opset
+CUSTOM_OPSET = onnxscript.values.Opset(domain="my_domain", version=1)
+
+@onnxscript.script(CUSTOM_OPSET)
+def my_complex_op(x: onnxscript.FLOAT, mask: onnxscript.BOOL) -> onnxscript.FLOAT:
+    """Custom operator that applies masking."""
+    return x * mask.cast(onnxscript.FLOAT)
+
+# Use in translation table
+custom_translation_table = {
+    my_pytorch_op: my_complex_op,
+}
+```
+
+Key considerations:
+- onnxscript functions use ONNX IR types (FLOAT, INT32, BOOL, etc.), not PyTorch types.
+- The function signature must match the number of inputs/outputs of the PyTorch operation.
+- Use `.cast()` for type conversions when needed.
+
+## Workflow
+
+### 0. Decide between direct onnxscript and hierarchical export
+
+Before writing code, ask:
+- **Is this a single atomic custom operation?** → Use direct onnxscript (Pattern 1).
+- **Is this a module hierarchy that needs reusable ONNX functions?** → Use hierarchical export (Pattern 2).
+- **Is this transformer attention?** → Use the built-in pattern (Pattern 3).
+
+If unsure, check:
+- [tests/expoter/test_dynamo_make_custom_op.py](../../../tests/expoter/test_dynamo_make_custom_op.py) for simple custom ops.
+- [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py) for hierarchical examples.
+- [tests/expoter/test_dynamo_build_onnxscript.py](../../../tests/expoter/test_dynamo_build_onnxscript.py) for build_onnxscript usage.
+
+### 1. Check torch version and dynamo availability
+
+```python
+import torch
+from inspect import signature
+
+# Verify dynamo parameter exists
+sig = signature(torch.onnx.export)
+if 'dynamo' not in sig.parameters:
+    # Use torch_export_handle_lower_version from HyperONNX
+    from hyperonnx.torch_export import torch_export_handle_lower_version
+    torch_export_handle_lower_version(model, args, f, dynamo=True, ...)
+else:
+    torch.onnx.export(model, args, f, dynamo=True, ...)
+```
+
+### 2. Identify and define custom operators
+
+For each custom operator:
+- Determine its input/output signatures.
+- Write an onnxscript function or use build_onnxscript.
+- Add it to custom_translation_table.
+
+Example:
+```python
+custom_translation_table = {}
+
+# Add manually-written onnxscript functions
+custom_translation_table[torch.ops.custom_lib.op1] = onnxscript_op1
+
+# Or add auto-generated ones from module specs
+custom_translation_table[MyModule] = build_onnxscript(my_module_spec)
+```
+
+### 3. Build custom_translation_table
+
+Use the context manager `replace_with_custom_op` for hierarchical export:
+
+```python
+from hyperonnx.exporter.dynamo import replace_with_custom_op
+
+with replace_with_custom_op(model, module_spec) as custom_translation_table:
+    torch.onnx.export(
+        model,
+        example_inputs,
+        'model.onnx',
+        dynamo=True,
+        custom_translation_table=custom_translation_table,
+        opset_version=18,
+    )
+```
+
+### 4. Call torch.onnx.export with dynamo=True
+
+```python
+torch.onnx.export(
+    model,
+    example_inputs,
+    'model.onnx',
+    input_names=['input'],
+    output_names=['output'],
+    opset_version=18,
+    dynamo=True,
+    custom_translation_table=custom_translation_table,
+    dynamic_shapes={'input': {0: 'batch_size'}},  # Optional: specify dynamic shapes
+)
+```
+
+Key parameters:
+- `dynamo=True`: Use torch.export instead of TorchScript.
+- `custom_translation_table`: Dict mapping callables to onnxscript functions.
+- `dynamic_shapes`: Preferred over dynamic_axes for dynamo path.
+- `opset_version`: Must align with onnxscript opset.
+
+### 5. Handle export failures and fallback
+
+The dynamo path may fail on complex models. HyperONNX provides fallback:
+
+```python
+from hyperonnx.hyper_export import export_hyper_onnx
+
+# export_hyper_onnx tries fallback automatically
+export_hyper_onnx(
+    model,
+    args,
+    'model.onnx',
+    dynamo=True,
+    # If dynamo export fails, HyperONNX may retry with dynamo=False
+)
+```
+
+To disable fallback and fail fast:
+```python
+# Call torch_export_handle_lower_version directly without fallback
+torch.onnx.export(
+    model,
+    args,
+    f,
+    dynamo=True,
+    # No fallback logic here
+)
+```
+
+### 6. Validate the exported ONNX model
+
+```python
+import onnx
+
+# Load and check
+onnx_model = onnx.load('model.onnx')
+onnx.checker.check_model(onnx_model, full_check=True)
+
+# Verify custom operators appear
+for node in onnx_model.graph.node:
+    print(f"{node.op_type} in domain {node.domain}")
+```
+
+## Common Errors & Recovery
+
+| Error | Cause | Fix |
+|-------|-------|-----|
+| `KeyError` in custom_translation_table | torch.ops key not in table | Verify the key exactly matches the PyTorch op reference |
+| `TypeError: ... got an unexpected keyword argument 'dynamo'` | torch.onnx.export too old | Use `torch_export_handle_lower_version` wrapper |
+| onnxscript function signature mismatch | ONNX function has wrong number of inputs/outputs | Verify spec['args'] and spec['output'] match onnxscript function |
+| `UnflatteningError` in build_onnxscript | Output structure not flattened | Use `plain_tensor_container` to extract only tensors |
+| Export succeeds but ONNX model is wrong | dynamic_shapes not specified | Add `dynamic_shapes` parameter for dynamic-shape models |
+| torch.export fails on data-dependent branches | Model has unhandled control flow | Wrap loops/conditionals with `torch.cond` or fall back to dynamo=False |
+
+## Branching Guidance
+
+If the request is about library implementation:
+- Prefer HyperONNX abstractions (export_hyper_onnx, replace_with_custom_op).
+- Preserve backward compatibility with older torch versions.
+- Use torch_export_handle_lower_version wrapper for version bridging.
+
+If the request is about custom operators:
+- Start with simple onnxscript functions (Pattern 1) for atomic ops.
+- Use build_onnxscript for auto-generation from module specs.
+- Write manual onnxscript only when fine-grained control is needed.
+
+If the request is about a failing export:
+- Check if the model uses data-dependent control flow; consider torch.cond.
+- Verify custom_translation_table keys match torch.ops exactly.
+- Check onnxscript function signatures match specs.
+- Try lowering torch version requirement (dynamo may not exist in older PyTorch).
+- Consider falling back to dynamo=False if dynamo export is unreliable for this model.
+
+## Quality Bar
+
+A good answer or code change using this skill should:
+- Clearly distinguish between dynamo=True (modern) and dynamo=False (legacy).
+- Use repository-native abstractions (export_hyper_onnx, replace_with_custom_op, build_onnxscript).
+- Specify dynamic_shapes when exporting models with variable batch size or sequence length.
+- Include custom_translation_table definitions with clear onnxscript signatures.
+- Provide validation (onnx.checker.check_model) and explain expected ONNX structure.
+- Mention torch version requirements and fallback behavior.
+
+## Example Prompts
+
+- Export this model using dynamo=True and custom attention operators.
+- Add dynamic shape support to this ONNX export.
+- Create a custom_translation_table for my hierarchical model.
+- Why does dynamo export fail on this model with torch.cond?
+- Migrate this export from dynamo=False to dynamo=True.
+- Build an onnxscript function for this custom module.
+
+## Weak Spots To Clarify Before Extending This Skill
+
+If you evolve this skill later, clarify these points first:
+
+- Which torch versions (2.5, 2.6, 2.8, 2.9+) should be officially supported for dynamo export examples?
+- Should we recommend torch.cond for all data-dependent branches, or only certain patterns?
+- Is build_onnxscript reliable enough to recommend for user-defined modules, or should we prefer manual onnxscript?
+- How to handle dtype inference in build_onnxscript for models with mixed precision (float16, bfloat16)?
+- Should fallback behavior be exposed as a user-facing option, or remain internal to HyperONNX?

hyperonnx-1.0.3/.claude/skills/torch-legacy-onnx-export/SKILL.md

@@ -0,0 +1,407 @@
+---
+name: torch-legacy-onnx-export
+description: Use this skill when working in HyperONNX and you need the legacy Torch ONNX export path with dynamo=False, especially for TorchScript-based export behavior, lower-version compatibility, custom attention registration, or post-export ONNX fixes.
+argument-hint: Describe the model, export entrypoint, torch version, and whether custom attention or custom input structures are involved.
+---
+
+# Torch Legacy ONNX Export
+
+Use this skill when the task is about the legacy Torch ONNX export path in HyperONNX, meaning the TorchScript-style path with dynamo=False.
+
+This skill is specific to this repository. It captures the actual workflow used by HyperONNX code and tests, rather than generic PyTorch advice.
+
+## Torch Behaviour
+
+- PyTorch 2.9 changed `torch.onnx.export(..., dynamo=...)` so that `dynamo=True` is the default.
+- On PyTorch 2.8 and earlier, do not assume the default selects the modern exporter. If legacy behavior matters, set `dynamo=False` explicitly. If modern exporter behavior matters, set `dynamo=True` explicitly.
+- In this repository, prefer being explicit either way. Do not rely on the library default because the default changed across PyTorch versions.
+
+⚠️ **Version Constraint**: This skill is tested and supported on **torch < 2.11** with **onnxscript < 0.6.0**. These versions have breaking changes in ONNX/module structure that require future work to resolve.
+
+## What This Skill Produces
+
+Produce one of these outcomes:
+
+- A correct export snippet that uses the legacy path with dynamo=False.
+- A HyperONNX code change that preserves legacy export compatibility.
+- A debugging checklist for why a legacy export is failing in this repository.
+- A validation plan that proves the exported ONNX model is structurally correct.
+
+## When To Use It
+
+Use this skill when the user asks for any of the following:
+
+- Use torch.onnx.export with the old or legacy path.
+- Force TorchScript export instead of dynamo export.
+- Keep HyperONNX compatible with older torch.onnx.export behavior.
+- Export a model with dynamo=False.
+- Make custom attention export work on the TorchScript path.
+- Understand why legacy export behaves differently from dynamo export.
+
+Do not use this skill for generic ONNX export advice that is unrelated to HyperONNX.
+
+## Repository Facts To Ground Decisions
+
+- HyperONNX exposes a version-bridging wrapper in [hyperonnx/torch_export.py](../../../hyperonnx/torch_export.py) named torch_export_handle_lower_version.
+- The main hierarchical export entrypoint is [hyperonnx/hyper_export.py](../../../hyperonnx/hyper_export.py), which accepts dynamo and may fall back between export modes.
+- The TorchScript branch for hierarchical replacement lives in [hyperonnx/exporter/__init__.py](../../../hyperonnx/exporter/__init__.py).
+- Transformer attention registration for the legacy path is implemented in [hyperonnx/transformers/attention.py](../../../hyperonnx/transformers/attention.py).
+- Repository examples already show dynamo=False in [README.md](../../../README.md) and [README_CN.md](../../../README_CN.md).
+- The duck-typing and symbolic operation system for TorchScript export is in [hyperonnx/exporter/torchscript.py](../../../hyperonnx/exporter/torchscript.py).
+
+## Registering Custom Operators for Legacy Export
+
+The legacy TorchScript export path requires explicit symbolic function registration for custom operators. Use this section when the task is to export custom PyTorch operations to ONNX while using dynamo=False.
+
+### Pattern 1: Direct symbolic registration (torch-native style)
+
+For simple custom operations that do not need hierarchical module wrapping:
+
+```python
+import torch
+from torch.onnx import register_custom_op_symbolic
+
+# Step 1: Define a symbolic function for ONNX
+def symbolic_my_op(g, input, param):
+    """ONNX symbolic for my_op.
+
+    Args:
+        g: The ONNX graph context.
+        input: The input tensor from PyTorch.
+        param: A parameter.
+
+    Returns:
+        ONNX operator node(s).
+    """
+    return g.op('custom_domain::MyOp', input, param_f=param)
+
+# Step 2: Register the symbolic before export
+register_custom_op_symbolic('custom_domain::my_op', symbolic_my_op, opset_version=18)
+
+# Step 3: Export with dynamo=False
+torch.onnx.export(
+    model,
+    example_inputs,
+    'model.onnx',
+    opset_version=18,
+    dynamo=False,
+)
+```
+
+Key rules:
+- The symbolic function name must match the registered name in `"<domain>::<op>"` format.
+- Always use opset_version consistently between registration and export.
+- Register before calling export; registration persists in the current process.
+
+### Pattern 2: HyperONNX hierarchical module with duck-typed custom op
+
+For complex models where the custom operation is itself a module that needs hierarchical export, use the duck-typing system from [hyperonnx/exporter/torchscript.py](../../../hyperonnx/exporter/torchscript.py):
+
+```python
+import torch
+from hyperonnx import export_hyper_onnx
+
+class MyCustomModule(torch.nn.Module):
+    """A module that implements custom logic."""
+    def forward(self, x):
+        # Custom computation
+        return x * 2 + 1
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(10, 20),
+    MyCustomModule(),
+    torch.nn.Linear(20, 5),
+)
+
+export_hyper_onnx(
+    model,
+    (torch.randn(1, 10),),
+    'model.onnx',
+    hiera=[MyCustomModule],  # Export MyCustomModule as a function
+    dynamo=False,
+    opset_version=18,
+)
+```
+
+How it works:
+1. When you pass `hiera=[MyCustomModule]`, HyperONNX will trace and export the module separately.
+2. The module is then replaced with a duck-type that acts like a custom operator node.
+3. The DuckForward class (from torchscript.py) handles the ONNX symbolic generation automatically.
+
+Reference: [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py) for hierarchical examples.
+
+### Pattern 3: Attention-specific registration
+
+For transformer attention layers, HyperONNX provides a pre-built pattern:
+
+```python
+import torch
+from hyperonnx.transformers.attention import register_attention_opsets
+
+# Register custom attention ONNX symbolic for the legacy path
+register_attention_opsets()
+
+# Then proceed with normal export
+torch.onnx.export(
+    model,
+    example_inputs,
+    'model.onnx',
+    opset_version=24,
+    dynamo=False,
+)
+```
+
+This registration only affects the TorchScript (legacy) path. See "Handle transformer attention correctly" in Workflow step 4 for post-export adjustments.
+
+## Understanding torchscript.py: Duck-Typing and Symbolic Export
+
+The [hyperonnx/exporter/torchscript.py](../../../hyperonnx/exporter/torchscript.py) module implements a clever pattern for exporting PyTorch modules as ONNX custom operators without requiring explicit symbolic functions. Here is how it works:
+
+### Core Concepts
+
+**DuckForward**: A torch.autograd.Function subclass that:
+- Stores the precomputed output in module_spec at forward time.
+- Implements a `symbolic()` static method that generates an ONNX node instead of computing.
+- Acts as a "duck type" – it quacks like a normal forward pass during tracing, but produces ONNX during export.
+
+**make_duck_forward**: Creates a DuckForward instance with closure over a specific module_spec.
+
+**replace_duck_forward**: Context manager that temporarily replaces module.forward with duck-typed versions.
+
+### Workflow (Three Phases)
+
+**Phase 1: Tracing and Caching**
+
+Before export, HyperONNX calls `trace_module_spec()` which:
+1. Registers forward hooks on all modules in the hiera list.
+2. Runs the model with example inputs to capture: input signatures, actual output values, loop count.
+3. Stores all this in module_spec.
+
+**Phase 2: Duck-Typing (During torch.onnx.export)**
+
+When `torch.onnx.export` traces the model with dynamo=False:
+1. `replace_duck_forward()` context manager replaces each hiera module's forward with a duck-typed version.
+2. Each duck forward is a DuckForward.apply() call that:
+   - During normal tracing: calls DuckForward.forward(), returns the precomputed cached output.
+   - During ONNX symbolic tracing: calls DuckForward.symbolic(), returns an ONNX op node.
+3. The exporter builds the ONNX graph, inserting custom op nodes where the duck forwards were.
+
+**Phase 3: Function Embedding**
+
+After export, HyperONNX uses other rewriters to:
+1. Embed the actual ONNX subgraph for each module into a function definition.
+2. Replace the custom op nodes with function calls.
+
+### Key Implementation Details
+
+**Forward method**: Stores precomputed output in module_spec, indexed by loop counter. Handles multiple calls to the same module and flattens nested output structures.
+
+**Symbolic method**: Uses the module's type_name (set during tracing) as the ONNX op name. Filters out None args and specifies the number of outputs.
+
+**Forward wrapper (_forward function)**:
+- Runs in inference mode to avoid backprop.
+- Flattens input structures the same way as the real forward.
+- Reorders kwargs to match the original function signature.
+- Calls DuckForward.apply which triggers either forward or symbolic depending on context.
+
+### Why This Pattern Matters for Legacy Export
+
+1. **No need for explicit symbolic functions**: The duck-typing pattern generates ONNX nodes automatically.
+2. **Handles complex module interactions**: Multiple calls, nested outputs, custom input types are all cached.
+3. **Integrates with TorchScript tracer**: The exporter sees normal forward calls; duck-typing is transparent.
+4. **Preserves correctness**: By caching actual outputs during tracing, the ONNX graph matches PyTorch semantics.
+
+### When to Use vs. When Not To Use
+
+**Use duck-typing (torchscript.py) when:**
+- You are using HyperONNX hierarchical export (passing hiera list).
+- The module is complex enough that writing a correct symbolic function is error-prone.
+- You want the output to be an embedded ONNX function, not just a custom op node.
+
+**Use direct symbolic registration when:**
+- You only have a single custom operator, not a module hierarchy.
+- You already have a working symbolic function.
+- You want the ONNX output to reference an external custom op domain.
+
+### Common Gotchas
+
+1. **Mismatched signatures**: If the cached output from tracing does not match the actual forward call structure, duck forward will return wrong data. Always verify with trace_module_spec.
+2. **Module called multiple times**: Duck forward increments loops and indexes into loop_outputs. Ensure tracing and export see the same number of calls.
+3. **Non-tensor outputs**: The symbolic method must return an op or tuple of ops. Use plain_tensor_container to extract just tensors.
+4. **Backward incompatibility**: Do not call backward on a DuckForward result. The backward method raises RuntimeError by design.
+
+## Workflow
+
+### 0. Decide between direct symbolic registration and duck-typing
+
+Before writing code, ask:
+- **Is this a single custom op that will appear in the final ONNX?** → Use direct symbolic registration (Pattern 1).
+- **Is this a PyTorch module that should become a reusable ONNX function?** → Use HyperONNX hierarchical export with duck-typing (Pattern 2).
+- **Is this transformer attention?** → Use the built-in pattern (Pattern 3) and refer to step 4 of Workflow.
+
+If unsure, check the test files:
+- [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py) for hierarchical examples.
+- [hyperonnx/transformers/attention.py](../../../hyperonnx/transformers/attention.py) for attention-specific patterns.
+
+### 1. Identify the export entrypoint
+
+First determine which layer of the repository the task belongs to:
+
+- If the task is about HyperONNX library internals, prefer torch_export_handle_lower_version instead of calling raw torch.onnx.export directly.
+- If the task is about exporting a full hierarchical model, work through export_hyper_onnx in [hyperonnx/hyper_export.py](../../../hyperonnx/hyper_export.py).
+- If the task is a direct usage example or a focused test, raw torch.onnx.export is acceptable when it matches the repository’s current test style.
+
+### 2. Choose the export mode deliberately
+
+Use this decision logic:
+
+- If the requirement is legacy or TorchScript export, set dynamo=False explicitly.
+- If the user only says export must work and does not require legacy behavior, note that HyperONNX may try fallback modes internally in some paths.
+- If the code is in trace_module_spec, remember the repository intentionally uses dynamo=False there to keep tracing behavior aligned with TorchScript export semantics.
+
+### 3. Apply the correct calling pattern
+
+For repository code changes:
+
+- Preserve the wrapper call shape used in [hyperonnx/torch_export.py](../../../hyperonnx/torch_export.py).
+- Keep kwargs handling compatible with torch versions below 2.5.0, where kwargs are merged into args.
+- Do not remove the OperatorExportTypes.ONNX_ATEN_FALLBACK behavior for older torch versions unless the task explicitly requires it and the change is validated.
+
+For direct usage examples:
+
+- Pass dynamo=False explicitly.
+- Provide input_names and output_names when the model interface is non-trivial.
+- Keep opset_version explicit when export semantics matter.
+
+Minimal example:
+
+```python
+import torch
+
+model = MyModule().eval()
+example_inputs = (x, y)
+
+torch.onnx.export(
+    model,
+    example_inputs,
+    "model.onnx",
+    input_names=["x", "y"],
+    output_names=["out"],
+    opset_version=19,
+    dynamo=False,
+)
+```
+
+### 4. Handle transformer attention correctly
+
+If the task involves HyperONNX transformer attention, the legacy path has a required branch:
+
+- Before export on the TorchScript path, call register_attention_opsets() from [hyperonnx/transformers/attention.py](../../../hyperonnx/transformers/attention.py).
+- This registration exists specifically for the torchscript export path with dynamo=False.
+- If the export also relies on the custom Attention op becoming default-domain ONNX Attention opset 24, promote the exported model with promote_onnx_model_to_opset24() after export.
+
+Pattern:
+
+```python
+import onnx
+import torch
+
+from hyperonnx.transformers.attention import (
+    promote_onnx_model_to_opset24,
+    register_attention_opsets,
+)
+
+register_attention_opsets()
+
+torch.onnx.export(
+    model,
+    example_inputs,
+    output_path,
+    opset_version=24,
+    dynamo=False,
+)
+
+onnx_model = onnx.load_model(str(output_path))
+onnx_model = promote_onnx_model_to_opset24(onnx_model)
+onnx.checker.check_model(onnx_model, full_check=True)
+```
+
+Reference implementation: [tests/transformers/test_attention.py](../../../tests/transformers/test_attention.py)
+
+### 5. Watch for input-structure limitations
+
+Legacy export does not make arbitrary Python objects exportable.
+
+Use these rules from the repository tests:
+
+- Plain custom objects fail export.
+- Tuple or list subclasses may pass flattening but still lose class behavior.
+- A torch.Tensor subclass can sometimes preserve the needed behavior well enough for export.
+
+Reference: [tests/expoter/test_non_pod_export.py](../../../tests/expoter/test_non_pod_export.py)
+
+If the user is trying to export non-POD inputs, explain that the fix is usually to flatten or tensorize inputs, not to keep arbitrary Python objects in the ONNX boundary.
+
+### 6. Validate the exported model
+
+Minimum completion checks:
+
+- The export artifact exists or the in-memory export object is non-null.
+- onnx.checker.check_model passes.
+- The chosen export path really used dynamo=False when legacy behavior was required.
+- If custom attention is involved, verify the expected Attention node and domain/opset normalization.
+- If changing library code, verify at least the smallest relevant test for the touched path.
+
+Good repository-aligned validation targets:
+
+- [tests/transformers/test_attention.py](../../../tests/transformers/test_attention.py)
+- [tests/transformers/test_static_cache.py](../../../tests/transformers/test_static_cache.py)
+- [tests/expoter/test_dynamo_replace_custom_op.py](../../../tests/expoter/test_dynamo_replace_custom_op.py) for contrast with the dynamo path
+
+## Branching Guidance
+
+If the request is about library implementation:
+
+- Prefer HyperONNX wrappers and existing abstractions.
+- Preserve backward-compatibility branches.
+- Avoid replacing the wrapper with direct raw torch.onnx.export unless there is a strong reason.
+
+If the request is about examples or docs:
+
+- Show the shortest correct snippet with dynamo=False.
+- Mention any required registration step up front.
+- Link the example back to the repository’s existing test or README pattern.
+
+If the request is about a failing export:
+
+- Check whether the failure is caused by unsupported input structures.
+- Check whether the user forgot dynamo=False.
+- Check whether custom attention registration was skipped.
+- Check whether post-export opset promotion is missing.
+- Check whether the code path should have used torch_export_handle_lower_version.
+
+## Quality Bar
+
+A good answer or code change using this skill should:
+
+- Distinguish clearly between legacy TorchScript export and dynamo export.
+- Use repository-native entrypoints and helpers.
+- Mention version-compatibility handling when touching core export code.
+- Include concrete validation, not just export code.
+- Avoid generic ONNX advice that ignores HyperONNX’s actual implementation.
+
+## Example Prompts
+
+- Use the legacy torch.onnx.export path for this HyperONNX test.
+- Add a sample that exports with dynamo=False and validates the ONNX.
+- Debug why this custom attention export fails on the TorchScript path.
+- Update HyperONNX export code without breaking lower torch versions.
+- Show me the correct legacy export flow for a transformer module in this repo.
+
+## Weak Spots To Clarify Before Extending This Skill
+
+If you evolve this skill later, clarify these points first:
+
+- Whether the team wants to prefer legacy export in docs, or only keep it as a compatibility path.
+- Which torch version range should be treated as officially supported for legacy export examples.
+- Whether future transformer examples should standardize on post-export opset promotion or hide that behind a helper.

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/.gitignore

@@ -1,5 +1,6 @@
 .*/
-!.github/
+!/.claude/
+!/.github/
 *.egg-info
 __pycache__
 # pdm build tool
@@ -14,3 +15,8 @@ tests/models/
 
 uv.lock
 .coverage
+
+# agents
+/.claude/settings.json
+/.github/agents/
+/.github/skills/

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyperonnx
-Version: 1.0.2
+Version: 1.0.3
 Summary: Copyright (C) 2025 The HYPERONNX Authors.
 Author-email: Wenyi Tang <wenyitang@outlook.com>
 Requires-Python: >=3.11
@@ -9,7 +9,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
 License-File: LICENSE
-Requires-Dist: onnxifier>=2.0.1
+Requires-Dist: onnxifier>=2.0.3
 Requires-Dist: onnxscript<0.6.0
 Requires-Dist: onnx>=1.18.0
 Requires-Dist: torch<2.11 ; extra == "cpu"

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/hyperonnx/__init__.py

@@ -19,4 +19,4 @@ from .hyper_export import export_hyper_onnx
 from .utils import HYPER_DOMAIN
 
 __all__ = ["auto_trace_method", "export_hyper_onnx", "HYPER_DOMAIN"]
-__version__ = "1.0.2"
+__version__ = "1.0.3"

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/hyperonnx/hyper_export.py

@@ -209,6 +209,8 @@ def _export_hiera(
     dynamo: bool,
     external_data: bool,
     do_optimization: bool,
+    fold_nodes_to_functions: bool,
+    fuse_constants_to_function: bool,
     external_directory: str | PathLike | None,
     module_spec: dict[Module, ModuleSpec],
     hiera: Collection[type[Module]],
@@ -244,6 +246,8 @@ def _export_hiera(
                 hiera=hiera,
                 module_spec=child_spec,
                 do_optimization=do_optimization,
+                fold_nodes_to_functions=fold_nodes_to_functions,
+                fuse_constants_to_function=fuse_constants_to_function,
                 external_directory=external_directory,
             )
             # update since child_spec may be updated
@@ -318,6 +322,8 @@ def export_hyper_onnx(  # noqa: C901
     hiera: Collection[type[Module]] | None = None,
     module_spec: dict[Module, ModuleSpec] | None = None,
     do_optimization: bool = True,
+    fold_nodes_to_functions: bool = True,
+    fuse_constants_to_function: bool = True,
     external_directory: str | PathLike | None = None,
     **_: Any,  # ignored options
 ) -> Any | None:
@@ -347,6 +353,10 @@ def export_hyper_onnx(  # noqa: C901
         module_spec (Optional[Dict[Module, ModuleSpec]]): A dictionary to store the
             detail spec of modules.
         do_optimization: Whether to optimize the exported ONNX model.
+        fold_nodes_to_functions: Whether to compose individual nodes into a function.
+            Defaults to True.
+        fuse_constants_to_function: Whether to fuse constants into function.
+            Defaults to True.
         external_directory (Optional[str | PathLike]): The directory to save the onnx
             model exported to be composed. If not specified, the model will be saved
             in the memory. Set to True if functions to be composed are too large.
@@ -389,7 +399,8 @@ def export_hyper_onnx(  # noqa: C901
                 if _dyn == dynamo:
                     logger.warning(
                         f"  Failed to export {model_typename}, "
-                        f"try use dynamo={not _dyn}"
+                        f"try use dynamo={not _dyn}, "
+                        "set ONNXIFIER_LOG_LEVEL=DEBUG to see error logs."
                     )
                 logger.debug(f"  <<<\n{ex}")
         raise RuntimeError(f"  Failed to export {model_typename}.")
@@ -411,6 +422,8 @@ def export_hyper_onnx(  # noqa: C901
         dynamo=dynamo,
         external_data=external_data,
         do_optimization=do_optimization,
+        fold_nodes_to_functions=fold_nodes_to_functions,
+        fuse_constants_to_function=fuse_constants_to_function,
         external_directory=external_directory,
         module_spec=module_spec,
         hiera=hiera,
@@ -460,7 +473,9 @@ def export_hyper_onnx(  # noqa: C901
     passes = [ComposeOnnxAsFunctionRewriter(HYPER_DOMAIN, tuple(module_spec.values()))]
     if do_optimization:
         passes.extend(OPTIMIZER_PASSES)
+    if fold_nodes_to_functions:
         passes.append(ComposeNodesToFunctionsRewriter(model_typename))
+    if fuse_constants_to_function:
         passes.append(FuseConstantsToFunctionRewriter())
     with chdir(external_directory):
         graph = PassManager(

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/hyperonnx/transformers/cache.py

@@ -33,6 +33,8 @@ class StaticCache(Tensor):
         We should treat DynamicCache as static for onnx export.
     """
 
+    # pylint: disable=unused-argument
+
     def update(
         self,
         key_states: Tensor,
@@ -61,6 +63,12 @@ class StaticCache(Tensor):
 
         layer_kv = self[layer_idx]
         keys, values = layer_kv.unbind()
+        if torch.onnx.is_in_onnx_export():
+            # We don't want to update cache because it's inplace. During exporting,
+            # the forward could be called multiple times and the cache will be falsely
+            # updated, so we skip cache update in onnx exporting phase.
+            keys = keys.clone()
+            values = values.clone()
         # Some old models give None for `cache_position` or even omit passing
         # `cache_kwargs` when used as cross-attention, in which case we should copy
         # the whole Layer (key_states.shape[-2] == self.max_cache_len)
@@ -82,6 +90,40 @@ class StaticCache(Tensor):
 
         return keys, values
 
+    def update_conv_state(
+        self, conv_states: torch.Tensor, layer_idx: int, **kwargs
+    ) -> torch.Tensor:
+        """
+        Updates the cache with the new `conv_states` for the layer `layer_idx`.
+
+        Parameters:
+            conv_states (`torch.Tensor`):
+                The new conv states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+
+        Return:
+            `torch.Tensor`: The updated conv states.
+        """
+        raise NotImplementedError("Qwen GDN conv state is not supported yet.")
+
+    def update_recurrent_state(
+        self, recurrent_states: torch.Tensor, layer_idx: int, **kwargs
+    ) -> torch.Tensor:
+        """
+        Updates the cache with the new `recurrent_states` for the layer `layer_idx`.
+
+        Parameters:
+            smm_states (`torch.Tensor`):
+                The new ssm states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+
+        Return:
+            `torch.Tensor`: The updated ssm states.
+        """
+        raise NotImplementedError("Qwen GDN recurrent state is not supported yet.")
+
     def early_initialization(
         self,
         batch_size: int,
@@ -107,11 +149,9 @@ class StaticCache(Tensor):
         """Returns the sequence length of the cache for the given layer."""
         if layer_idx >= self.shape[0]:
             return 0
-        return int(self[0, 0].any(dim=-1).sum().item())
+        return int(self[layer_idx, 0].any(dim=-1).sum().item())
 
-    def get_mask_sizes(
-        self, cache_position: torch.Tensor, layer_idx: int
-    ) -> tuple[int, int]:
+    def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]:
         """
         Return a tuple (kv_length, kv_offset) corresponding to the length and offset
         that will be returned for the given layer at `layer_idx`. The masks are then
@@ -119,8 +159,9 @@ class StaticCache(Tensor):
         for each layer.
         """
         if layer_idx >= self.shape[0]:
-            return cache_position.shape[0], 0
-        return self.shape[4], 0
+            if isinstance(query_length, torch.Tensor) and query_length.ndim > 0:
+                return query_length.shape[0], 0
+        return self.max_cache_len, 0
 
     def get_max_cache_shape(self, layer_idx: int = 0) -> int:
         """
@@ -140,6 +181,18 @@ class StaticCache(Tensor):
         if self.get_seq_length() > 0:
             self.copy_(self.index_select(0, beam_idx.to(self.device)))
 
+    def crop(self, max_length: int):
+        """Do nothing for static cache."""
+        return
+
+    def batch_repeat_interleave(self, repeats: int):
+        """Do nothing for static cache."""
+        return
+
+    def batch_select_indices(self, indices: torch.Tensor):
+        """Do nothing for static cache."""
+        return
+
     @property
     def max_batch_size(self) -> int:
         """Return the maximum batch size of the cache"""

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/hyperonnx/transformers/patch.py

@@ -1,5 +1,5 @@
 """
-Copyright (C) 2025 The HYPERONNX Authors.
+Copyright (C) 2026 The HYPERONNX Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@ import importlib
 from contextlib import contextmanager
 from unittest.mock import patch
 
+import torch
 from packaging.version import Version
 
 
@@ -34,6 +35,7 @@ def patch_transformers():
     TRANSFORMERS_HAS_SDPA_ATTENTION = Version("4.48.0")
     TRANSFORMERS_HAS_SDPA_MASK = Version("4.53.0")
     TRANSFORMERS_FIX_VMAP = Version("5.0.0")
+    TRANSFORMERS_DEPRECATE_CACHE_POSITION = Version("5.4.0")
     CURR_VER = Version(transformers.__version__)
 
     patches = []
@@ -56,6 +58,33 @@ def patch_transformers():
             "transformers.integrations.sdpa_attention.use_gqa_in_sdpa", use_gqa_in_sdpa
         )
         patches.append(p2)
+    if CURR_VER >= TRANSFORMERS_DEPRECATE_CACHE_POSITION:
+        original_sdpa_mask = transformers.masking_utils.sdpa_mask
+
+        def patch_sdpa_mask(
+            batch_size: int,
+            q_length: int,
+            kv_length: int,
+            *args,
+            **kwargs,
+        ):
+            # In onnx exporting (torch.onnx.is_in_onnx_export()==True), the shape of a
+            # Tensor is still a tensor:
+            # type(tensor([1]).shape[0]) == torch.Tensor
+            # So the wrap in sdpa_mask is wrong in this case. We forcely convert the
+            # shape to integer.
+            def _to_int(x):
+                if isinstance(x, torch.Tensor) and x.ndim == 0:
+                    return int(x.item())
+                return x
+
+            q_length = _to_int(q_length)
+            kv_length = _to_int(kv_length)
+
+            return original_sdpa_mask(batch_size, q_length, kv_length, *args, **kwargs)
+
+        p3 = patch("transformers.masking_utils.sdpa_mask", patch_sdpa_mask)
+        patches.append(p3)
 
     try:
         for p in patches:

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
     "Topic :: Software Development",
 ]
 dynamic = ["version", "description"]
-dependencies = ["onnxifier>=2.0.1", "onnxscript<0.6.0", "onnx>=1.18.0"]
+dependencies = ["onnxifier>=2.0.3", "onnxscript<0.6.0", "onnx>=1.18.0"]
 
 [project.urls]
 Home = "https://github.com/LoSealL/hyperonnx"
@@ -47,7 +47,7 @@ dev = [
     "pyright>=1.1.408",
     "pytest",
     "pytest-cov",
-    "transformers>=4.48.0,<5.4.0",
+    "transformers>=4.48.0",
 ]
 
 [tool.setuptools]
@@ -56,7 +56,7 @@ packages = ["hyperonnx"]
 
 [tool.pyright]
 include = ["hyperonnx"]
-exclude = ["tests"]
+exclude = ["**/node_modules", "**/__pycache__", "**/.*", ".venv", "tests"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

{hyperonnx-1.0.2 → hyperonnx-1.0.3}/tests/test_export_hyper_onnx.py

@@ -140,6 +140,7 @@ def test_export_resnet_with_basicblock_and_bottleneck(dynamo):
             input_names=["input.1"],
             hiera=[ResNet, BasicBlock, Bottleneck],
             do_optimization=False,
+            fold_nodes_to_functions=False,
             dynamo=dynamo,
         )
         model = onnx.load_from_string(f.getvalue())
@@ -229,7 +230,6 @@ def test_export_llama_transformers(dynamo, cache, tmp_path):
         input_ids=torch.randint(config.vocab_size, [1, 16]).long(),
         attention_mask=torch.ones([1, 16]),
         position_ids=torch.arange(16).long()[None],
-        cache_position=torch.arange(16).long(),
     )
     if cache:
         # example_inputs["use_cache"] = True
@@ -243,10 +243,12 @@ def test_export_llama_transformers(dynamo, cache, tmp_path):
         )
     with torch.no_grad():
         ref_output = llama(**example_inputs)
+        if cache:
+            example_inputs["past_key_values"].fill_(0)
     with BytesIO() as f, patch_transformers(), torch.no_grad():
-        input_names = ["input_ids", "attention_mask", "position_ids", "cache_position"]
+        input_names = ["input_ids", "attention_mask", "position_ids"]
         if cache:
-            input_names.insert(-1, "past_key_values")
+            input_names.append("past_key_values")
         output_names = ["last_hidden_state"]
         if cache:
             output_names.append("present_key_values")
@@ -311,6 +313,8 @@ def test_export_rnnlike(dynamo):
             hiera=[RNNCell],
             input_names=["x"],
             do_optimization=False,
+            fold_nodes_to_functions=False,
+            fuse_constants_to_function=False,
             dynamo=dynamo,
         )
         onnx_model = onnx.load_from_string(f.getvalue())