Repository: pytorch
Issue: [pt2] Unable to trace LSTM with dynamic sequence length
Labels: triaged,oncall: pt2,module: dynamic shapes,export-triaged,oncall: export

### 🐛 Unable to trace LSTM with dynamic batch size

Export/PT2(?) is unable to trace LSTM with a dynamic batch size because this [line](https://github.com/pytorch/pytorch/blob/6317a0350e5f960fd9a56714f2602fdcf4e35c39/torch/_decomp/decompositions.py#L2979-L2981) in the decomposition loops over the batch dimension, causing a specialization in the batch dimension. This happens when MKLDNN is not enabled.

A solution might be to replace this loop with a `scan` operator?

Repro:
```python
    def test_lstm(self):
        import torch
        from torch._export import capture_pre_autograd_graph, dynamic_dim

        class MyModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.lstm = torch.nn.LSTM(input_size=4, hidden_size=5, num_layers=1)

            def forward(self, inputs: torch.Tensor) -> torch.Tensor:
                return self.lstm(inputs)

        inp = (torch.ones(3, 4),)
        torch._C._set_mkldnn_enabled(False)
        gm = capture_pre_autograd_graph(MyModule(), inp, constraints=[dynamic_dim(inp[0], 0)])
        gm.print_readable()
```

Example output:
```
(pytorch10) angelayi@devgpu022 /data/users/angelayi/pytorch [main] $ TORCH_LOGS=dynamo,dynamic python test/export/test_export.py -k test_lstm
[2023-12-04 12:50:11,246] [0/0] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo start tracing forward /data/users/angelayi/pytorch/test/export/test_export.py:626
[2023-12-04 12:50:11,247] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] create_env
[2023-12-04 12:50:11,966] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] create_symbol s0 = 3 for L['inputs'].size()[0] [2, 9223372036854775806]
[2023-12-04 12:50:22,556] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] set_replacement s0 = 3
[2023-12-04 12:50:22,556] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] set_replacement s0 = 3
[2023-12-04 12:50:22,560] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] eval Eq(s0, 3) [guard added] at nn/modules/module.py:1520 in _call_impl (_ops.py:541 in decompose)
[2023-12-04 12:50:22,610] [0/0] torch._dynamo.symbolic_convert: [INFO] Step 1: torchdynamo done tracing forward (RETURN_VALUE)
[2023-12-04 12:50:22,613] [0/0] torch._dynamo.output_graph: [INFO] Step 2: calling compiler function dynamo_normalization_capturing_compiler
[2023-12-04 12:50:22,613] [0/0] torch._dynamo.output_graph: [INFO] Step 2: done compiler function dynamo_normalization_capturing_compiler
[2023-12-04 12:50:22,616] [0/0] torch.fx.experimental.symbolic_shapes: [INFO] produce_guards
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR] Error while creating guard:
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR] Name: ''
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Source: shape_env
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Create Function: SHAPE_ENV
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Guard Types: None
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Code List: None
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Object Weakref: None
[2023-12-04 12:50:22,616] [0/0] torch._guards: [ERROR]     Guarded Class Weakref: None
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR] Created at:
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]   File "/data/users/angelayi/pytorch/torch/_dynamo/convert_frame.py", line 509, in transform
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]     tracer = InstructionTranslator(
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]   File "/data/users/angelayi/pytorch/torch/_dynamo/symbolic_convert.py", line 2061, in __init__
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]     output=OutputGraph(
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]   File "/data/users/angelayi/pytorch/torch/_dynamo/output_graph.py", line 296, in __init__
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]     self.init_ambient_guards()
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]   File "/data/users/angelayi/pytorch/torch/_dynamo/output_graph.py", line 366, in init_ambient_guards
[2023-12-04 12:50:22,617] [0/0] torch._guards: [ERROR]     self.guards.add(ShapeEnvSource().make_guard(GuardBuilder.SHAPE_ENV))
E
======================================================================
ERROR: test_lstm (__main__.TestExport)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/angelayi/pytorch/torch/testing/_internal/common_utils.py", line 2652, in wrapper
    method(*args, **kwargs)
  File "/data/users/angelayi/pytorch/test/export/test_export.py", line 631, in test_lstm
    gm = capture_pre_autograd_graph(MyModule(), inp, constraints=[dynamic_dim(inp[0], 0)])
  File "/data/users/angelayi/pytorch/torch/_export/__init__.py", line 372, in capture_pre_autograd_graph
    m = b(
  File "/data/users/angelayi/pytorch/torch/_dynamo/eval_frame.py", line 1382, in inner
    raise constraint_violation_error
  File "/data/users/angelayi/pytorch/torch/_dynamo/eval_frame.py", line 1342, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/_dynamo/eval_frame.py", line 655, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/data/users/angelayi/pytorch/torch/_dynamo/convert_frame.py", line 383, in _convert_frame_assert
    compiled_product = _compile(
  File "/data/users/angelayi/pytorch/torch/_dynamo/convert_frame.py", line 646, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/data/users/angelayi/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/angelayi/pytorch/torch/_dynamo/convert_frame.py", line 626, in compile_inner
    check_fn = CheckFunctionManager(
  File "/data/users/angelayi/pytorch/torch/_dynamo/guards.py", line 1011, in __init__
    guard.create(builder)
  File "/data/users/angelayi/pytorch/torch/_guards.py", line 246, in create
    return self.create_fn(builder, self)
  File "/data/users/angelayi/pytorch/torch/_dynamo/guards.py", line 670, in SHAPE_ENV
    guards = output_graph.shape_env.produce_guards(
  File "/data/users/angelayi/pytorch/torch/fx/experimental/symbolic_shapes.py", line 2939, in produce_guards
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (L['inputs'].size()[0])! For more information, run with TORCH_LOGS=dynamic.
  - Not all values of L['inputs'].size()[0] = L['inputs'].size()[0] in the specified range are valid because L['inputs'].size()[0] was inferred to be a constant (3).
```

### Versions

main

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @suo @ydwu4 @sxu @tarun292