Optimize

scripts/compare_onnx_tflite.py

@@ -0,0 +1,667 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Script to compare the results of an ONNX model with a TFLite model given the same input.
+Optionally also compare with Tract runtime for ONNX.
+Created by Copilot.
+
+Usage:
+    python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite
+    python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --input input.npy
+    python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --rtol 1e-5 --atol 1e-5
+    python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --benchmark
+    python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --use-tract --benchmark
+"""
+
+import argparse
+import time
+import numpy as np
+import onnxruntime as ort
+import tensorflow as tf
+from typing import Dict, List, Tuple, Optional, Any
+
+try:
+    import tract
+
+    TRACT_AVAILABLE = True
+except ImportError:
+    TRACT_AVAILABLE = False
+
+
+def load_onnx_model(onnx_path: str) -> ort.InferenceSession:
+    """Load an ONNX model and return an inference session."""
+    print(f"Loading ONNX model from: {onnx_path}")
+    session = ort.InferenceSession(onnx_path)
+    return session
+
+
+def load_tflite_model(tflite_path: str) -> tf.lite.Interpreter:
+    """Load a TFLite model and return an interpreter."""
+    print(f"Loading TFLite model from: {tflite_path}")
+    interpreter = tf.lite.Interpreter(model_path=tflite_path)
+    interpreter.allocate_tensors()
+    return interpreter
+
+
+def load_tract_model(onnx_path: str) -> Optional[Any]:
+    """Load an ONNX model using tract and return a runnable model."""
+    if not TRACT_AVAILABLE:
+        print("Tract is not available. Install with: pip install tract")
+        return None
+    print(f"Loading ONNX model with tract from: {onnx_path}")
+    model = tract.onnx().model_for_path(onnx_path).into_optimized().into_runnable()
+    return model
+
+
+def get_onnx_model_info(session: ort.InferenceSession) -> Tuple[List, List]:
+    """Get input and output information from ONNX model."""
+    inputs = session.get_inputs()
+    outputs = session.get_outputs()
+
+    print("\nONNX Model Information:")
+    print("Inputs:")
+    for inp in inputs:
+        print(f"  - Name: {inp.name}, Shape: {inp.shape}, Type: {inp.type}")
+    print("Outputs:")
+    for out in outputs:
+        print(f"  - Name: {out.name}, Shape: {out.shape}, Type: {out.type}")
+
+    return inputs, outputs
+
+
+def get_tflite_model_info(interpreter: tf.lite.Interpreter) -> Tuple[List, List]:
+    """Get input and output information from TFLite model."""
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+
+    print("\nTFLite Model Information:")
+    print("Inputs:")
+    for inp in input_details:
+        print(f"  - Name: {inp['name']}, Shape: {inp['shape']}, Type: {inp['dtype']}")
+    print("Outputs:")
+    for out in output_details:
+        print(f"  - Name: {out['name']}, Shape: {out['shape']}, Type: {out['dtype']}")
+
+    return input_details, output_details
+
+
+def generate_random_inputs(onnx_inputs: List, seed: int = 42) -> Dict[str, np.ndarray]:
+    """Generate random inputs based on ONNX model input specs."""
+    np.random.seed(seed)
+    inputs = {}
+
+    print("\nGenerating random inputs:")
+    for inp in onnx_inputs:
+        # Handle dynamic dimensions
+        shape = []
+        for dim in inp.shape:
+            if isinstance(dim, str) or dim is None or dim < 0:
+                # Default to 1 for dynamic dimensions
+                shape.append(1)
+            else:
+                shape.append(dim)
+
+        # Generate random data based on type
+        if "float" in inp.type.lower():
+            data = np.random.randn(*shape).astype(np.float32)
+        elif "int64" in inp.type.lower():
+            data = np.random.randint(0, 100, size=shape).astype(np.int64)
+        elif "int32" in inp.type.lower():
+            data = np.random.randint(0, 100, size=shape).astype(np.int32)
+        else:
+            # Default to float32
+            data = np.random.randn(*shape).astype(np.float32)
+
+        inputs[inp.name] = data
+        print(f"  - {inp.name}: shape={data.shape}, dtype={data.dtype}")
+
+    return inputs
+
+
+def load_inputs_from_file(input_path: str) -> Dict[str, np.ndarray]:
+    """Load inputs from a numpy file (.npy or .npz)."""
+    print(f"\nLoading inputs from: {input_path}")
+
+    if input_path.endswith(".npz"):
+        data = np.load(input_path)
+        inputs = {key: data[key] for key in data.files}
+    elif input_path.endswith(".npy"):
+        data = np.load(input_path)
+        # Assume single input
+        inputs = {"input": data}
+    else:
+        raise ValueError("Input file must be .npy or .npz format")
+
+    for name, value in inputs.items():
+        print(f"  - {name}: shape={value.shape}, dtype={value.dtype}")
+
+    return inputs
+
+
+def run_onnx_model(
+    session: ort.InferenceSession, inputs: Dict[str, np.ndarray]
+) -> List[np.ndarray]:
+    """Run inference on ONNX model."""
+    print("\nRunning ONNX model inference...")
+    outputs = session.run(None, inputs)
+    return outputs
+
+
+def run_tflite_model(
+    interpreter: tf.lite.Interpreter, inputs: Dict[str, np.ndarray], input_details: List
+) -> List[np.ndarray]:
+    """Run inference on TFLite model."""
+    print("Running TFLite model inference...")
+
+    # Set input tensors
+    for i, detail in enumerate(input_details):
+        # Try to match by name or use order
+        input_data = None
+        if detail["name"] in inputs:
+            input_data = inputs[detail["name"]]
+        elif len(inputs) == 1:
+            # If only one input, use it
+            input_data = list(inputs.values())[0]
+        elif i < len(inputs):
+            # Use by order
+            input_data = list(inputs.values())[i]
+        else:
+            raise ValueError(f"Cannot match input for TFLite input {detail['name']}")
+
+        # Ensure correct dtype
+        if input_data.dtype != detail["dtype"]:
+            input_data = input_data.astype(detail["dtype"])
+
+        interpreter.set_tensor(detail["index"], input_data)
+
+    # Run inference
+    interpreter.invoke()
+
+    # Get output tensors
+    output_details = interpreter.get_output_details()
+    outputs = []
+    for detail in output_details:
+        outputs.append(interpreter.get_tensor(detail["index"]))
+
+    return outputs
+
+
+def run_tract_model(model: Any, inputs: Dict[str, np.ndarray]) -> List[np.ndarray]:
+    """Run inference on tract model."""
+    if model is None:
+        return []
+    print("Running tract model inference...")
+
+    # Convert inputs to list (tract expects a list of tensors)
+    input_list = list(inputs.values())
+
+    # Run inference
+    outputs = model.run(input_list)
+
+    # Convert outputs to numpy arrays
+    result = []
+    for output in outputs:
+        result.append(output.to_numpy())
+
+    return result
+
+
+def benchmark_onnx_model(
+    session: ort.InferenceSession,
+    inputs: Dict[str, np.ndarray],
+    num_runs: int = 100,
+    warmup_runs: int = 10,
+) -> Dict[str, float]:
+    """Benchmark ONNX model inference speed."""
+    print(f"\nBenchmarking ONNX model ({warmup_runs} warmup + {num_runs} test runs)...")
+
+    # Warmup runs
+    for _ in range(warmup_runs):
+        session.run(None, inputs)
+
+    # Timed runs
+    times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        session.run(None, inputs)
+        end = time.perf_counter()
+        times.append((end - start) * 1000)  # Convert to ms
+
+    return {
+        "mean": np.mean(times),
+        "median": np.median(times),
+        "std": np.std(times),
+        "min": np.min(times),
+        "max": np.max(times),
+    }
+
+
+def benchmark_tflite_model(
+    interpreter: tf.lite.Interpreter,
+    inputs: Dict[str, np.ndarray],
+    input_details: List,
+    num_runs: int = 100,
+    warmup_runs: int = 10,
+) -> Dict[str, float]:
+    """Benchmark TFLite model inference speed."""
+    print(f"Benchmarking TFLite model ({warmup_runs} warmup + {num_runs} test runs)...")
+
+    # Prepare inputs
+    def set_inputs():
+        for i, detail in enumerate(input_details):
+            input_data = None
+            if detail["name"] in inputs:
+                input_data = inputs[detail["name"]]
+            elif len(inputs) == 1:
+                input_data = list(inputs.values())[0]
+            elif i < len(inputs):
+                input_data = list(inputs.values())[i]
+            else:
+                raise ValueError(
+                    f"Cannot match input for TFLite input {detail['name']}"
+                )
+
+            if input_data.dtype != detail["dtype"]:
+                input_data = input_data.astype(detail["dtype"])
+
+            interpreter.set_tensor(detail["index"], input_data)
+
+    # Warmup runs
+    for _ in range(warmup_runs):
+        set_inputs()
+        interpreter.invoke()
+
+    # Timed runs
+    times = []
+    for _ in range(num_runs):
+        set_inputs()
+        start = time.perf_counter()
+        interpreter.invoke()
+        end = time.perf_counter()
+        times.append((end - start) * 1000)  # Convert to ms
+
+    return {
+        "mean": np.mean(times),
+        "median": np.median(times),
+        "std": np.std(times),
+        "min": np.min(times),
+        "max": np.max(times),
+    }
+
+
+def benchmark_tract_model(
+    model: Any,
+    inputs: Dict[str, np.ndarray],
+    num_runs: int = 100,
+    warmup_runs: int = 10,
+) -> Optional[Dict[str, float]]:
+    """Benchmark tract model inference speed."""
+    if model is None:
+        return None
+    print(f"Benchmarking tract model ({warmup_runs} warmup + {num_runs} test runs)...")
+
+    # Convert inputs to list
+    input_list = list(inputs.values())
+
+    # Warmup runs
+    for _ in range(warmup_runs):
+        model.run(input_list)
+
+    # Timed runs
+    times = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        model.run(input_list)
+        end = time.perf_counter()
+        times.append((end - start) * 1000)  # Convert to ms
+
+    return {
+        "mean": np.mean(times),
+        "median": np.median(times),
+        "std": np.std(times),
+        "min": np.min(times),
+        "max": np.max(times),
+    }
+
+
+def print_benchmark_results(
+    onnx_stats: Dict[str, float],
+    tflite_stats: Dict[str, float],
+    tract_stats: Optional[Dict[str, float]] = None,
+) -> None:
+    """Print benchmark comparison results."""
+    print("\n" + "=" * 80)
+    print("BENCHMARK RESULTS")
+    print("=" * 80)
+
+    print("\nONNX Model:")
+    print(f"  Mean:   {onnx_stats['mean']:.3f} ms")
+    print(f"  Median: {onnx_stats['median']:.3f} ms")
+    print(f"  Std:    {onnx_stats['std']:.3f} ms")
+    print(f"  Min:    {onnx_stats['min']:.3f} ms")
+    print(f"  Max:    {onnx_stats['max']:.3f} ms")
+
+    print("\nTFLite Model:")
+    print(f"  Mean:   {tflite_stats['mean']:.3f} ms")
+    print(f"  Median: {tflite_stats['median']:.3f} ms")
+    print(f"  Std:    {tflite_stats['std']:.3f} ms")
+    print(f"  Min:    {tflite_stats['min']:.3f} ms")
+    print(f"  Max:    {tflite_stats['max']:.3f} ms")
+
+    if tract_stats:
+        print("\nTract Model:")
+        print(f"  Mean:   {tract_stats['mean']:.3f} ms")
+        print(f"  Median: {tract_stats['median']:.3f} ms")
+        print(f"  Std:    {tract_stats['std']:.3f} ms")
+        print(f"  Min:    {tract_stats['min']:.3f} ms")
+        print(f"  Max:    {tract_stats['max']:.3f} ms")
+
+    print("\nComparison:")
+    speedup = tflite_stats["mean"] / onnx_stats["mean"]
+    if speedup > 1:
+        print(f"  ONNX Runtime is {speedup:.2f}x faster than TFLite")
+    else:
+        print(f"  TFLite is {1 / speedup:.2f}x faster than ONNX Runtime")
+    print(f"  Difference: {abs(onnx_stats['mean'] - tflite_stats['mean']):.3f} ms")
+
+    if tract_stats:
+        speedup_tract = tflite_stats["mean"] / tract_stats["mean"]
+        if speedup_tract > 1:
+            print(f"  Tract is {speedup_tract:.2f}x faster than TFLite")
+        else:
+            print(f"  TFLite is {1 / speedup_tract:.2f}x faster than Tract")
+        print(f"  Difference: {abs(tract_stats['mean'] - tflite_stats['mean']):.3f} ms")
+
+        speedup_ort = onnx_stats["mean"] / tract_stats["mean"]
+        if speedup_ort > 1:
+            print(f"  Tract is {speedup_ort:.2f}x faster than ONNX Runtime")
+        else:
+            print(f"  ONNX Runtime is {1 / speedup_ort:.2f}x faster than Tract")
+        print(f"  Difference: {abs(tract_stats['mean'] - onnx_stats['mean']):.3f} ms")
+
+    print("=" * 80)
+
+
+def compare_outputs(
+    onnx_outputs: List[np.ndarray],
+    tflite_outputs: List[np.ndarray],
+    tract_outputs: Optional[List[np.ndarray]] = None,
+    rtol: float = 1e-5,
+    atol: float = 1e-5,
+) -> bool:
+    """Compare outputs from ONNX, TFLite, and optionally Tract models."""
+    print("\n" + "=" * 80)
+    print("COMPARISON RESULTS")
+    print("=" * 80)
+
+    if len(onnx_outputs) != len(tflite_outputs):
+        print(
+            f"❌ Number of outputs differs: ONNX={len(onnx_outputs)}, TFLite={len(tflite_outputs)}"
+        )
+        return False
+
+    if tract_outputs and len(onnx_outputs) != len(tract_outputs):
+        print(
+            f"❌ Number of outputs differs: ONNX={len(onnx_outputs)}, Tract={len(tract_outputs)}"
+        )
+        return False
+
+    all_match = True
+    for i, (onnx_out, tflite_out) in enumerate(zip(onnx_outputs, tflite_outputs)):
+        tract_out = tract_outputs[i] if tract_outputs else None
+
+        print(f"\nOutput {i}:")
+        print(f"  ONNX Runtime shape: {onnx_out.shape}, dtype: {onnx_out.dtype}")
+        print(f"  TFLite shape:       {tflite_out.shape}, dtype: {tflite_out.dtype}")
+        if tract_out is not None:
+            print(f"  Tract shape:        {tract_out.shape}, dtype: {tract_out.dtype}")
+
+        if onnx_out.shape != tflite_out.shape:
+            print("  ❌ Shape mismatch between ONNX and TFLite!")
+            all_match = False
+            continue
+
+        if tract_out is not None and onnx_out.shape != tract_out.shape:
+            print("  ❌ Shape mismatch between ONNX and Tract!")
+            all_match = False
+            continue
+
+        # Convert to same dtype for comparison
+        if onnx_out.dtype != tflite_out.dtype:
+            print("  ⚠️  Different dtypes, converting to float32 for comparison")
+            onnx_out = onnx_out.astype(np.float32)
+            tflite_out = tflite_out.astype(np.float32)
+
+        if tract_out is not None and onnx_out.dtype != tract_out.dtype:
+            tract_out = tract_out.astype(np.float32)
+
+        # Compute statistics - ONNX vs TFLite
+        print("\n  ONNX Runtime vs TFLite:")
+        diff = np.abs(onnx_out - tflite_out)
+        max_diff = np.max(diff)
+        mean_diff = np.mean(diff)
+        is_close = np.allclose(onnx_out, tflite_out, rtol=rtol, atol=atol)
+
+        print(f"    Max difference:  {max_diff:.10f}")
+        print(f"    Mean difference: {mean_diff:.10f}")
+        print(f"    Relative tolerance: {rtol}")
+        print(f"    Absolute tolerance: {atol}")
+
+        if is_close:
+            print("    ✅ Outputs match within tolerance")
+        else:
+            print("    ❌ Outputs do NOT match within tolerance")
+            all_match = False
+
+            # Show some sample values
+            print("\n    Sample values (first 5 elements):")
+            flat_onnx = onnx_out.flatten()[:5]
+            flat_tflite = tflite_out.flatten()[:5]
+            for j, (o, t) in enumerate(zip(flat_onnx, flat_tflite)):
+                print(
+                    f"      [{j}] ONNX: {o:.10f}, TFLite: {t:.10f}, Diff: {abs(o - t):.10f}"
+                )
+
+        # Compute statistics - ONNX vs Tract
+        if tract_out is not None:
+            print("\n  ONNX Runtime vs Tract:")
+            diff_tract = np.abs(onnx_out - tract_out)
+            max_diff_tract = np.max(diff_tract)
+            mean_diff_tract = np.mean(diff_tract)
+            is_close_tract = np.allclose(onnx_out, tract_out, rtol=rtol, atol=atol)
+
+            print(f"    Max difference:  {max_diff_tract:.10f}")
+            print(f"    Mean difference: {mean_diff_tract:.10f}")
+
+            if is_close_tract:
+                print("    ✅ Outputs match within tolerance")
+            else:
+                print("    ❌ Outputs do NOT match within tolerance")
+                all_match = False
+
+                # Show some sample values
+                print("\n    Sample values (first 5 elements):")
+                flat_onnx_tract = onnx_out.flatten()[:5]
+                flat_tract = tract_out.flatten()[:5]
+                for j, (o, tr) in enumerate(zip(flat_onnx_tract, flat_tract)):
+                    print(
+                        f"      [{j}] ONNX: {o:.10f}, Tract: {tr:.10f}, Diff: {abs(o - tr):.10f}"
+                    )
+
+            # Compute statistics - TFLite vs Tract
+            print("\n  TFLite vs Tract:")
+            diff_tflite_tract = np.abs(tflite_out - tract_out)
+            max_diff_tflite_tract = np.max(diff_tflite_tract)
+            mean_diff_tflite_tract = np.mean(diff_tflite_tract)
+            is_close_tflite_tract = np.allclose(
+                tflite_out, tract_out, rtol=rtol, atol=atol
+            )
+
+            print(f"    Max difference:  {max_diff_tflite_tract:.10f}")
+            print(f"    Mean difference: {mean_diff_tflite_tract:.10f}")
+
+            if is_close_tflite_tract:
+                print("    ✅ Outputs match within tolerance")
+            else:
+                print("    ❌ Outputs do NOT match within tolerance")
+                all_match = False
+
+    print("\n" + "=" * 80)
+    if all_match:
+        print("✅ ALL OUTPUTS MATCH!")
+    else:
+        print("❌ SOME OUTPUTS DO NOT MATCH")
+    print("=" * 80)
+
+    return all_match
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compare ONNX and TFLite model outputs",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Compare with random inputs
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite
+
+  # Compare with custom inputs from file
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --input input.npz
+
+  # Compare with custom tolerances
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --rtol 1e-3 --atol 1e-3
+
+  # Save outputs for inspection
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --save-outputs
+
+  # Benchmark execution speed
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --benchmark
+
+  # Benchmark with custom number of runs
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --benchmark --num-runs 200 --warmup-runs 20
+
+  # Compare with tract runtime as well
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --use-tract
+
+  # Benchmark all three runtimes
+  python compare_onnx_tflite.py --onnx model.onnx --tflite model.tflite --use-tract --benchmark
+        """,
+    )
+
+    parser.add_argument("--onnx", required=True, help="Path to ONNX model")
+    parser.add_argument("--tflite", required=True, help="Path to TFLite model")
+    parser.add_argument("--input", help="Path to input file (.npy or .npz)")
+    parser.add_argument(
+        "--rtol", type=float, default=1e-5, help="Relative tolerance (default: 1e-5)"
+    )
+    parser.add_argument(
+        "--atol", type=float, default=1e-5, help="Absolute tolerance (default: 1e-5)"
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for input generation (default: 42)",
+    )
+    parser.add_argument(
+        "--save-outputs", action="store_true", help="Save outputs to files"
+    )
+    parser.add_argument(
+        "--benchmark",
+        action="store_true",
+        help="Benchmark execution speed of both models",
+    )
+    parser.add_argument(
+        "--num-runs",
+        type=int,
+        default=100,
+        help="Number of benchmark runs (default: 100)",
+    )
+    parser.add_argument(
+        "--warmup-runs",
+        type=int,
+        default=10,
+        help="Number of warmup runs (default: 10)",
+    )
+    parser.add_argument(
+        "--use-tract", action="store_true", help="Also test with tract ONNX runtime"
+    )
+
+    args = parser.parse_args()
+
+    # Load models
+    onnx_session = load_onnx_model(args.onnx)
+    tflite_interpreter = load_tflite_model(args.tflite)
+
+    # Load tract model if requested
+    tract_model = None
+    if args.use_tract:
+        if not TRACT_AVAILABLE:
+            print(
+                "\n⚠️  Warning: Tract is not installed. Install with: pip install tract"
+            )
+            print("Continuing without tract comparison...\n")
+        else:
+            tract_model = load_tract_model(args.onnx)
+
+    # Get model info
+    onnx_inputs, onnx_outputs = get_onnx_model_info(onnx_session)
+    tflite_input_details, tflite_output_details = get_tflite_model_info(
+        tflite_interpreter
+    )
+
+    # Prepare inputs
+    if args.input:
+        inputs = load_inputs_from_file(args.input)
+    else:
+        inputs = generate_random_inputs(onnx_inputs, seed=args.seed)
+
+    # Run inference
+    onnx_results = run_onnx_model(onnx_session, inputs)
+    tflite_results = run_tflite_model(tflite_interpreter, inputs, tflite_input_details)
+    tract_results = None
+    if tract_model:
+        tract_results = run_tract_model(tract_model, inputs)
+
+    # Save outputs if requested
+    if args.save_outputs:
+        print("\nSaving outputs...")
+        np.savez("onnx_outputs.npz", *onnx_results)
+        np.savez("tflite_outputs.npz", *tflite_results)
+        print("  - onnx_outputs.npz")
+        print("  - tflite_outputs.npz")
+        if tract_results:
+            np.savez("tract_outputs.npz", *tract_results)
+            print("  - tract_outputs.npz")
+
+    # Compare results
+    match = compare_outputs(
+        onnx_results, tflite_results, tract_results, rtol=args.rtol, atol=args.atol
+    )
+
+    # Benchmark if requested
+    if args.benchmark:
+        onnx_stats = benchmark_onnx_model(
+            onnx_session, inputs, args.num_runs, args.warmup_runs
+        )
+        tflite_stats = benchmark_tflite_model(
+            tflite_interpreter,
+            inputs,
+            tflite_input_details,
+            args.num_runs,
+            args.warmup_runs,
+        )
+        tract_stats = None
+        if tract_model:
+            tract_stats = benchmark_tract_model(
+                tract_model, inputs, args.num_runs, args.warmup_runs
+            )
+        print_benchmark_results(onnx_stats, tflite_stats, tract_stats)
+
+    # Return exit code
+    return 0 if match else 1
+
+
+if __name__ == "__main__":
+    exit(main())

scripts/optimize.py

@@ -24,13 +24,13 @@ class ReplaceDftWithMatMulRule(onnxscript.rewriter.RewriteRuleClassBase):
         num_freqs = dft_size // 2 + 1
 
         # Vectorized creation of DFT matrix
+        n = np.arange(dft_size, dtype=np.float32)[:, np.newaxis]  # Shape: (dft_size, 1)
         k = np.arange(num_freqs, dtype=np.float32)[
-            :, np.newaxis
-        ]  # Shape: (num_freqs, 1)
-        n = np.arange(dft_size, dtype=np.float32)[np.newaxis, :]  # Shape: (1, dft_size)
+            np.newaxis, :
+        ]  # Shape: (1, num_freqs)
         dft_matrix = np.cos(
             2 * np.pi * k * n / dft_size
-        )  # Shape: (num_freqs, dft_size)
+        )  # Shape: (dft_size, num_freqs)
 
         # Create constant node for the DFT matrix
         dft_matrix = op.initializer(ir.tensor(dft_matrix), name=f"{x.name}_dft_matrix")
@@ -41,13 +41,75 @@ class ReplaceDftWithMatMulRule(onnxscript.rewriter.RewriteRuleClassBase):
         return result
 
 
+class ReplaceSplit(onnxscript.rewriter.RewriteRuleClassBase):
+    def pattern(self, op, x):
+        return op.Split(x, _allow_other_inputs=True, _outputs=["split_out_1", "split_out_2"])
+    
+    def rewrite(self, op, x: ir.Value, **kwargs):
+        zero = op.initializer(ir.tensor(np.array([0], dtype=np.int64)), "zero")
+        batch_size = op.Gather(x, zero)
+        sample_size = op.initializer(ir.tensor(np.array([144000], dtype=np.int32)), "sample_size")
+        return batch_size, sample_size
+
+
+class RemoveCast(onnxscript.rewriter.RewriteRuleClassBase):
+    def pattern(self, op, x):
+        return op.Cast(x)
+
+    def rewrite(self, op, x: ir.Value, **kwargs):
+        return op.Identity(x)
+
 model = ir.load("model.onnx")
+
+# Set dynamic axes
+model.graph.inputs[0].shape = ir.Shape(["batch", 144000])
+model.graph.outputs[0].shape = ir.Shape(["batch", 6522])
+
+onnxscript.rewriter.rewrite(model, [ReplaceDftWithMatMulRule().rule(), ReplaceSplit().rule(), RemoveCast().rule()])
+
+# Change all int32 initializers to int64
+initializers = list(model.graph.initializers.values())
+for initializer in initializers:
+    if initializer.dtype == ir.DataType.INT32:
+        int32_array = initializer.const_value.numpy()
+        int64_array = int32_array.astype(np.int64)
+        new_initializer = ir.val(initializer.name, const_value=ir.tensor(int64_array))
+        model.graph.initializers.pop(initializer.name)
+        model.graph.initializers.add(new_initializer)
+        initializer.replace_all_uses_with(new_initializer)
+
 onnxscript.optimizer.optimize(
     model, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
 )
-onnxscript.rewriter.rewrite(model, [ReplaceDftWithMatMulRule().rule()])
+
+# Remove Slice-Reshape
+def remove_slice_reshape(model: ir.Model):
+    mul_node = model.graph.node("model/MEL_SPEC1/Mul")
+    first_reshape = model.graph.node("model/MEL_SPEC1/stft/frame/Reshape_1")
+    first_shape = ir.val("first_shape", const_value=ir.tensor([-1, 72000, 2], dtype=ir.DataType.INT64))
+    model.graph.initializers.add(first_shape)
+    second_reshape = model.graph.node("model/MEL_SPEC2/stft/frame/Reshape_1")
+    second_shape = ir.val("second_shape", const_value=ir.tensor([-1, 18000, 8], dtype=ir.DataType.INT64))
+    model.graph.initializers.add(second_shape)
+
+    # Replace with Mul-Reshape-Gather
+    first_reshape.replace_input_with(0, mul_node.outputs[0])
+    first_reshape.replace_input_with(1, first_shape)
+    second_reshape.replace_input_with(0, first_reshape.outputs[0])
+    second_reshape.replace_input_with(1, second_shape)
+
+
+
+remove_slice_reshape(model)
+# Run DCE again
+onnxscript.optimizer.optimize(
+    model, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
+)
+
 
 onnx_ir.passes.common.ClearMetadataAndDocStringPass()(model)
 model.ir_version = 10
+model.producer_name = "onnx-ir"
+model.graph.name = "BirdNET-v2.4"
 
 ir.save(model, "birdnet.onnx")