Repository: pytorch
Issue: Error in torch.export for torch.ops.aten.chunk for dynamic shape
Labels: triaged,oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

While exporting a  model with `torch.ops.aten.chunk` it decomposes into` torch.ops.aten.split.Tensor` in `torch.export`. With dynamic inputs the below code-
```
import torch
import torch_tensorrt
class TestChunk(torch.nn.Module):
    def forward(self, input):
        out = torch.ops.aten.chunk.default(input, 3, 0)
        return out

inputs = [torch.randn(3)]
inputs_zero_shape = torch.export.Dim("shape", min=1, max=3)
dynamic_shapes = [[torch.export.Dim("shape", min=1, max=3)]]
exp_program = torch.export.export(TestChunk(), tuple(inputs), dynamic_shapes=dynamic_shapes)
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs)
# Run inference
trt_gm(*inputs)
```
Fails in torch.export. The error is the following
```
Traceback (most recent call last):
  File "/code/torch_trt/TensorRT/tests/py/dynamo/conversion/split_dynamic.py", line 14, in <module>
    exp_program = torch.export.export(TestChunk(), tuple(inputs), dynamic_shapes=dynamic_shapes)
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/__init__.py", line 172, in export
    return _export(
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1013, in wrapper
    raise e
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 986, in wrapper
    ep = fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/exported_program.py", line 97, in wrapper
    return fn(*args, **kwargs)
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1921, in _export
    export_artifact = export_func(  # type: ignore[operator]
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1220, in _strict_export
    return _strict_export_lower_to_aten_ir(
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 1248, in _strict_export_lower_to_aten_ir
    gm_torch_level = _export_to_torch_ir(
  File "/root/.pyenv/versions/3.10.14/lib/python3.10/site-packages/torch/export/_trace.py", line 572, in _export_to_torch_ir
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (shape)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of shape = L['input'].size()[0] in the specified range shape <= 3 satisfy the generated guard Eq(((L['input'].size()[0] + ((L['input'].size()[0] + 2)//3) - 1)//(((L['input'].size()[0] + 2)//3))), 3).
Specializations unexpectedly required (shape)! For more information, run with TORCH_LOGS="+dynamic".
  - solving the guards generated for shape = L['input'].size()[0] resulted in a specialized value of 3.

Suggested fixes:
  shape = 3
```
It seems to fail in the torch.export guard. Could someone please look into this?

### Versions

Torch- 2.5.0.dev20240827+cu124

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4