[Copilot]

LSTM and GRU exports compute output shapes at runtime after every layer, adding kernels that cannot be constant-folded with dynamic batches. The reported four-layer LSTM workload incurs roughly 10% overhead.

• Simplify layer conversion: Share a helper across both lowerings. Use Squeeze(Y, [1]) for unidirectional outputs; use Transpose followed by Reshape([0, 0, 2 * hidden_size]) for bidirectional outputs.
• Preserve dynamic dimensions: Reshape zeros copy sequence and batch dimensions without shape arithmetic.
• Guard graph structure and correctness: Add multilayer regression coverage across directions, layouts, sequence lengths, and dynamic batch sizes.

The four-layer LSTM export drops from 104 to 68 nodes, eliminating all output-conversion Shape and Reshape nodes.

• Fixes [torchlib] aten_lstm / aten_gru: runtime shape computation between layers makes multi-layer RNNs ~10% slower #3096

[codecov]

❌ 4 Tests Failed:

Tests completed	Failed	Passed	Skipped
10023	4	10019	1854

View the top 2 failed test(s) by shortest run time

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__addbmm_cpu_float32

Stack Traces | 0.422s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 50 / 50 (100.0%)
E   Greatest absolute difference: nan at index (0, 0) (up to 2e-05 allowed)
E   Greatest relative difference: nan at index (0, 0) (up to 2e-05 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__addbmm_cpu_float16

Stack Traces | 0.456s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 50 / 50 (100.0%)
E   Greatest absolute difference: nan at index (0, 0) (up to 0.02 allowed)
E   Greatest relative difference: nan at index (0, 0) (up to 0.2 allowed)

View the full list of 8 ❄️ flaky test(s)

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1207_test_optional_get_element_optional_tensor

Flake rate in main: 100.00% (Passed 0 times, Failed 1879 times)

Stack Traces | 0.005s run time

..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
    code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
    return exporter.export(model_onnx, function_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
    add(self._import_onnx_types(proto))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
    ts = _translate_type(t.type)
         ^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
    raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E   NotImplementedError: Unable to translate type optional_type {
E     elem_type {
E       tensor_type {
E         elem_type: 1
E         shape {
E           dim {
E             dim_value: 4
E           }
E         }
E       }
E     }
E   }
E    into onnxscript type.

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1211_test_optional_has_element_tensor_input

Flake rate in main: 100.00% (Passed 0 times, Failed 1879 times)

Stack Traces | 0.008s run time

..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
    code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
    return exporter.export(model_onnx, function_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
    add(self._import_onnx_types(proto))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
    ts = _translate_type(t.type)
         ^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
    raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E   NotImplementedError: Unable to translate type optional_type {
E     elem_type {
E       tensor_type {
E         elem_type: 1
E         shape {
E           dim {
E             dim_value: 4
E           }
E         }
E       }
E     }
E   }
E    into onnxscript type.

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1212_test_optional_has_element_optional_input

Flake rate in main: 100.00% (Passed 0 times, Failed 1879 times)

Stack Traces | 0.006s run time

..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
    code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
    return exporter.export(model_onnx, function_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
    add(self._import_onnx_types(proto))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
    ts = _translate_type(t.type)
         ^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
    raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E   NotImplementedError: Unable to translate type optional_type {
E     elem_type {
E       tensor_type {
E         elem_type: 1
E         shape {
E           dim {
E             dim_value: 4
E           }
E         }
E       }
E     }
E   }
E    into onnxscript type.

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1215_test_optional_has_element_empty_optional_input

Flake rate in main: 100.00% (Passed 0 times, Failed 1879 times)

Stack Traces | 0.004s run time

..../test_torch_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export_test.py:260: in test_export2python_produces_correct_onnx_script_model
    code = onnx_export.export2python(
onnxscript/backend/onnx_export.py:919: in export2python
    return exporter.export(model_onnx, function_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:821: in export
    add(self._import_onnx_types(proto))
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:798: in _import_onnx_types
    ts = _translate_type(t.type)
         ^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/backend/onnx_export.py:124: in _translate_type
    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/onnx_types.py:269: in onnx_type_to_onnxscript_repr
    raise NotImplementedError(f"Unable to translate type {onnx_type!r} into onnxscript type.")
E   NotImplementedError: Unable to translate type optional_type {
E     elem_type {
E       tensor_type {
E         elem_type: 6
E         shape {
E         }
E       }
E     }
E   }
E    into onnxscript type.

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__baddbmm_cpu_float16

Flake rate in main: 8.02% (Passed 860 times, Failed 75 times)

Stack Traces | 0.471s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 250 / 250 (100.0%)
E   Greatest absolute difference: nan at index (0, 0, 0) (up to 0.01 allowed)
E   Greatest relative difference: nan at index (0, 0, 0) (up to 0.001 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__baddbmm_cpu_float32

Flake rate in main: 14.53% (Passed 594 times, Failed 101 times)

Stack Traces | 0.406s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 250 / 250 (100.0%)
E   Greatest absolute difference: nan at index (0, 0, 0) (up to 1e-05 allowed)
E   Greatest relative difference: nan at index (0, 0, 0) (up to 1.3e-06 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float16

Flake rate in main: 8.02% (Passed 860 times, Failed 75 times)

Stack Traces | 0.481s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 83 / 125 (66.4%)
E   Greatest absolute difference: 0.8115234375 at index (0, 0, 0) (up to 0.0007 allowed)
E   Greatest relative difference: 1.9990234375 at index (0, 0, 0) (up to 0.1 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logit_cpu_float32

Flake rate in main: 8.02% (Passed 860 times, Failed 75 times)

Stack Traces | 0.363s run time

tests\function_libs\torch_lib\ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 81 / 125 (64.8%)
E   Greatest absolute difference: 0.8109303712844849 at index (0, 0, 0) (up to 1e-05 allowed)
E   Greatest relative difference: 2.0 at index (0, 0, 0) (up to 1.3e-06 allowed)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

🟢 Approval recommended

The conversions preserve dimension ordering and dynamic sizes, with no blocking issues identified in static review; tests were not executed.

No open issues 🎉

What changed in this PR

Simplifies LSTM and GRU output conversion to remove per-layer runtime shape computation while preserving dynamic batch dimensions.

Changes:

• Shares a conversion helper using Squeeze or constant-shape Reshape.
• Adds multilayer regression coverage for both directions, layouts, sequence lengths, and varying batch sizes.

File	Description
tests/​function_libs/​torch_lib/​e2e_ops_tests.py	Checks graph structure and numerical correctness across RNN configurations.
onnxscript/​function_libs/​torch_lib/​ops/​core.py	Replaces runtime shape arithmetic with simplified shared conversion.

_{Review effort: Balanced}