| 2767 | ) |
| 2768 | |
| 2769 | def replace_input(inp): |
| 2770 | assert isinstance(spec, InputSpec) |
| 2771 | if spec.type == "user_input": |
| 2772 | arg = spec.user_input.arg |
| 2773 | if arg.type == "as_tensor": |
| 2774 | t = arg.as_tensor |
| 2775 | t.name = replace_table[t.name] |
| 2776 | elif arg.type == "as_sym_int": |
| 2777 | s = arg.as_sym_int |
| 2778 | if s.type == "as_name": |
| 2779 | s.as_name = replace_table[s.as_name] |
| 2780 | elif s.type == "as_int": |
| 2781 | pass |
| 2782 | else: |
| 2783 | raise AssertionError(f"Unknown sym_int type: {s}") |
| 2784 | elif arg.type in ( |
| 2785 | "as_none", |
| 2786 | "as_bool", |
| 2787 | "as_int", |
| 2788 | "as_float", |
| 2789 | "as_string", |
| 2790 | "as_custom_obj", |
| 2791 | ): |
| 2792 | return |
| 2793 | else: |
| 2794 | raise AssertionError(f"Unknown input type: {arg}") |
| 2795 | elif spec.type == "parameter": |
| 2796 | t = spec.parameter.arg |
| 2797 | t.name = replace_table[t.name] |
| 2798 | elif spec.type == "buffer": |
| 2799 | t = spec.buffer.arg |
| 2800 | t.name = replace_table[t.name] |
| 2801 | elif spec.type == "tensor_constant": |
| 2802 | t = spec.tensor_constant.arg |
| 2803 | t.name = replace_table[t.name] |
| 2804 | elif spec.type == "custom_obj": |
| 2805 | return |
| 2806 | elif spec.type == "token": |
| 2807 | tok = spec.token.arg |
| 2808 | tok.name = replace_table[tok.name] |
| 2809 | elif spec.type == "constant_input": |
| 2810 | return |
| 2811 | else: |
| 2812 | raise AssertionError(f"Unknown input type: {spec}") |
| 2813 | |
| 2814 | def replace_output(out): |
| 2815 | assert isinstance(spec, OutputSpec) |