Repository: pytorch
Issue: dynamic tensor shape of vmap dimention
Labels: actionable,module: functorch,module: dynamic shapes,export-triaged,oncall: export

### 🐛 Describe the bug

I've have a model that internally consists of multiple data streams, each of which having their own unique parameters. To optimize for performance I had stacked different parameters for all the streams together and issued a single call for linear operations with help of `torch.vmap`. This approach worked fine back in torch 2.3, but ever since upgrading to torch-2.4 the compile operation fails if I mark the vmap dimension as dynamic. I've stripped all the details from the code down to the smallest repro script:

```python
import torch
import typing
import torch._dynamo.config

class StreamedLinear(torch.nn.Module):
	def __init__(self, n_streams : int, d_in : int, d_out : int, bias:bool = True):
		super().__init__()
		self.n_streams = n_streams
		self.weight_streams = torch.nn.Parameter(torch.empty((n_streams, d_out, d_in)))

		if bias:
			self.bias_streams = torch.nn.Parameter(torch.empty((n_streams, d_out)))
		else:
			self.bias_streams = None

		self.weight: typing.Optional[torch.Tensor] = None
		self.bias: typing.Optional[torch.Tensor] = None
		self.backend = torch.vmap(torch.nn.functional.linear)
		
	def forward(self, stream_idx:torch.Tensor, x:torch.Tensor):
		weight = self.weight_streams.index_select(0, stream_idx)
												 
		if self.bias is None:
			return self.backend(x, weight)
		else:
			bias = self.bias_streams.index_select(0, stream_idx)
			return self.backend(x, weight, bias)


def exec(model:torch.nn.Module , args : torch.Tensor):
	
	stream_selection = list(range(args.shape[0]))
	print(type(model))
	model.apply(lambda module: module.activate_streams(stream_selection) if hasattr(module, 'activate_streams') else None)
	
	return model(args)


if __name__ == '__main__':
	
	test_layer = StreamedLinear(2, 3, 4, False)
	test_args = torch.randn(2, 5, 3)
	stream_idx = torch.tensor([0, 1], dtype=int)

	stream_dim = torch.export.Dim('streams')
	compiled = torch.export.export(test_layer, (stream_idx, test_args), dynamic_shapes={'stream_idx':{0:stream_dim}, 'x':{0:stream_dim}})
	print(compiled)
```

It seems the constraints created because of vmap are enforcing constant guards rather than dynamic/relative ones. In my searches I've seen that similar issues have previously been reported but nothing was explicitly addressing vmap transforms. I additionally noted that in 3.3 this was issue was bypassed since the graph broke automatically when hitting function transformers, but the default behavior has changed in 3.4.

I have captured the following logs: Hopefully it helps:
```
V0801 01:28:54.031000 140445296705536 torch/fx/experimental/symbolic_shapes.py:2529] [0/0] create_env
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0] create_symbol s0 = 2 for L['stream_idx'].size()[0] [2, 9223372036854775806] (_dynamo/variables/builder.py:2276 in <lambda>)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0] For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0] Stack (most recent call last):
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/test_compile.py", line 47, in <module>
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     compiled = torch.export.export(test_layer, (stream_idx, test_args), dynamic_shapes={'stream_idx':{0:stream_dim}, 'x':{0:stream_dim}})
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 174, in export
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return _export(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 928, in wrapper
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     ep = fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 89, in wrapper
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1455, in _export
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     aten_export_artifact = export_func(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1060, in _strict_export
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     gm_torch_level = _export_to_torch_ir(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 512, in _export_to_torch_ir
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     gm_torch_level, _ = torch._dynamo.export(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1379, in inner
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     result_traced = opt_f(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return self._call_impl(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return forward_call(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 433, in _fn
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return self._call_impl(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return forward_call(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1116, in __call__
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return self._torchdynamo_orig_callable(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 472, in __call__
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return _compile(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 84, in wrapper_function
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return StrobelightCompileTimeProfiler.profile_compile_time(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return func(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/usr/lib/python3.10/contextlib.py", line 79, in inner
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return func(*args, **kwds)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 817, in _compile
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     guarded_code = compile_inner(code, one_graph, hooks, transform)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     r = func(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 636, in compile_inner
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     out_code = transform_code_object(code, transform)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1185, in transform_code_object
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     transformations(instructions, code_options)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 178, in _fn
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 564, in transform
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     tracer = InstructionTranslator(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2416, in __init__
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     self.symbolic_locals = variables.LazyVariableTracker.realize_all(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 120, in realize_all
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     result = {k: cls.realize_all(v, cache) for k, v in list(value.items())}
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 120, in <dictcomp>
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     result = {k: cls.realize_all(v, cache) for k, v in list(value.items())}
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 106, in realize_all
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     result = cls.realize_all(value.realize(), cache)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     self._cache.realize()
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     self.vt = VariableBuilder(tx, self.source)(self.value)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 314, in __call__
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     vt = self._wrap(value)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 456, in _wrap
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return type_dispatch(self, value)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1303, in wrap_tensor
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     tensor_variable = wrap_fx_proxy(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1713, in wrap_fx_proxy
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1820, in wrap_fx_proxy_cls
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     example_value = wrap_to_fake_tensor_and_record(example_value, tx=tx, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 2275, in wrap_to_fake_tensor_and_record
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     fake_e = wrap_fake_exception(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 1300, in wrap_fake_exception
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn()
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 2276, in <lambda>
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     lambda: tx.fake_mode.from_tensor(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1941, in from_tensor
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return self.fake_tensor_converter.from_real_tensor(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 332, in from_real_tensor
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     out = self.meta_converter(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 1621, in __call__
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     r = self.meta_tensor(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 1408, in meta_tensor
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     ) = sym_sizes_strides_storage_offset(t, source, symbolic_context)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_subclasses/meta_utils.py", line 738, in sym_sizes_strides_storage_offset
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return shape_env._create_symbolic_sizes_strides_storage_offset(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 245, in wrapper
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 3151, in _create_symbolic_sizes_strides_storage_offset
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     size: List[sympy.Expr] = self._produce_dyn_sizes_from_int_tuple(ex_size, source, symbolic_context)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 3027, in _produce_dyn_sizes_from_int_tuple
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     size.append(self.create_symbol(
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/recording.py", line 245, in wrapper
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     return fn(*args, **kwargs)
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 3549, in create_symbol
I0801 01:28:54.080000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0]     self.log.info(
V0801 01:28:54.081000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5167] [0/0] eval True == True [statically known]
V0801 01:28:54.081000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5167] [0/0] eval False == False [statically known]
I0801 01:28:54.087000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3549] [0/0] create_symbol s1 = 2 for L['x'].size()[0] [2, 9223372036854775806] (_dynamo/variables/builder.py:2276 in <lambda>), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s1"
V0801 01:28:54.130000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5167] [0/0] eval False == False [statically known]
I0801 01:28:54.135000 140445296705536 torch/fx/experimental/symbolic_shapes.py:4831] [0/0] set_replacement s1 = s0 (solve) VR[2, 9223372036854775806]
I0801 01:28:54.136000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5082] [0/0] eval Eq(s0, s1) [guard added] at _dynamo/polyfill.py:21 in any (_dynamo/variables/tensor.py:1041 in evaluate_expr), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s0, s1)"
V0801 01:28:54.144000 140445296705536 torch/fx/experimental/symbolic_shapes.py:4697] [0/0] _update_var_to_range s0 = VR[2, 2] (update)
I0801 01:28:54.146000 140445296705536 torch/fx/experimental/symbolic_shapes.py:4831] [0/0] set_replacement s0 = 2 (solve) VR[2, 2]
I0801 01:28:54.146000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5082] [0/0] eval Eq(s0, 2) [guard added] at _functorch/vmap.py:476 in _flat_vmap (_dynamo/variables/tensor.py:1041 in evaluate_expr), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED="Eq(s0, 2)"
V0801 01:28:54.160000 140445296705536 torch/fx/experimental/symbolic_shapes.py:4697] [0/0] _update_var_to_range s1 = VR[2, 2] (update)
I0801 01:28:54.161000 140445296705536 torch/fx/experimental/symbolic_shapes.py:4831] [0/0] set_replacement s1 = 2 (find) VR[2, 2]
V0801 01:28:54.170000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5167] [0/0] eval True == True [statically known]
V0801 01:28:54.182000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5156] [0/0] eval 2 [trivial]
I0801 01:28:54.205000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3639] [0/0] produce_guards
V0801 01:28:54.206000 140445296705536 torch/fx/experimental/symbolic_shapes.py:5167] [0/0] eval True == True [statically known]
V0801 01:28:54.206000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['stream_idx'].size()[0] 2 StrictMinMaxConstraint(warn_only=False, vr=VR[0, 9223372036854775806])
V0801 01:28:54.206000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['stream_idx'].stride()[0] 1 None
V0801 01:28:54.206000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['stream_idx'].storage_offset() 0 None
V0801 01:28:54.206000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].size()[0] 2 StrictMinMaxConstraint(warn_only=False, vr=VR[0, 9223372036854775806])
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].size()[1] 5 None
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].size()[2] 3 None
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].stride()[0] 15 None
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].stride()[1] 3 None
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].stride()[2] 1 None
V0801 01:28:54.207000 140445296705536 torch/fx/experimental/symbolic_shapes.py:3821] [0/0] track_symint L['x'].storage_offset() 0 None
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0] Error while creating guard:
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0] Name: ''
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Source: shape_env
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Create Function: SHAPE_ENV
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Guard Types: None
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Code List: None
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Object Weakref: None
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     Guarded Class Weakref: None
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0] Traceback (most recent call last):
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_guards.py", line 260, in create
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     return self.create_fn(builder, self)
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 1717, in SHAPE_ENV
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     guards = output_graph.shape_env.produce_guards(
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 4163, in produce_guards
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]     raise ConstraintViolationError(
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0] torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (streams)! For more information, run with TORCH_LOGS="+dynamic".
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]   - Not all values of streams = L['stream_idx'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).
E0801 01:28:54.208000 140445296705536 torch/_guards.py:262] [0/0]   - Not all values of streams = L['x'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0] Created at:
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 564, in transform
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]     tracer = InstructionTranslator(
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2360, in __init__
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]     output=OutputGraph(
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 313, in __init__
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]     self.init_ambient_guards()
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]   File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 452, in init_ambient_guards
E0801 01:28:54.209000 140445296705536 torch/_guards.py:264] [0/0]     self.guards.add(ShapeEnvSource().make_guard(GuardBuilder.SHAPE_ENV))
Traceback (most recent call last):
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 512, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1425, in inner
    raise constraint_violation_error
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1379, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1116, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 726, in compile_inner
    check_fn = CheckFunctionManager(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 2130, in __init__
    guard.create(builder)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_guards.py", line 260, in create
    return self.create_fn(builder, self)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/_dynamo/guards.py", line 1717, in SHAPE_ENV
    guards = output_graph.shape_env.produce_guards(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/fx/experimental/symbolic_shapes.py", line 4163, in produce_guards
    raise ConstraintViolationError(
torch.fx.experimental.symbolic_shapes.ConstraintViolationError: Constraints violated (streams)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of streams = L['stream_idx'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).
  - Not all values of streams = L['x'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).

Suggested fixes:
  streams = 2

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ali/test_compile.py", line 47, in <module>
    compiled = torch.export.export(test_layer, (stream_idx, test_args), dynamic_shapes={'stream_idx':{0:stream_dim}, 'x':{0:stream_dim}})
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/__init__.py", line 174, in export
    return _export(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 945, in wrapper
    raise e
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 928, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/exported_program.py", line 89, in wrapper
    return fn(*args, **kwargs)
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1455, in _export
    aten_export_artifact = export_func(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 1060, in _strict_export
    gm_torch_level = _export_to_torch_ir(
  File "/home/ali/.local/lib/python3.10/site-packages/torch/export/_trace.py", line 529, in _export_to_torch_ir
    raise UserError(UserErrorType.CONSTRAINT_VIOLATION, str(e))  # noqa: B904
torch._dynamo.exc.UserError: Constraints violated (streams)! For more information, run with TORCH_LOGS="+dynamic".
  - Not all values of streams = L['stream_idx'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).
  - Not all values of streams = L['x'].size()[0] in the specified range are valid because streams was inferred to be a constant (2).

Suggested fixes:
  streams = 2
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats constrain_symbol_range: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats evaluate_expr: CacheInfo(hits=11, misses=8, maxsize=256, currsize=8)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats _simplify_floor_div: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats _maybe_guard_rel: CacheInfo(hits=0, misses=2, maxsize=256, currsize=2)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats _find: CacheInfo(hits=19, misses=7, maxsize=None, currsize=2)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats has_hint: CacheInfo(hits=0, misses=0, maxsize=256, currsize=0)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats size_hint: CacheInfo(hits=1, misses=2, maxsize=256, currsize=2)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats simplify: CacheInfo(hits=2, misses=9, maxsize=None, currsize=3)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats _update_divisible: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V0801 01:28:54.230000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats replace: CacheInfo(hits=196, misses=34, maxsize=None, currsize=9)
V0801 01:28:54.231000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats _maybe_evaluate_static: CacheInfo(hits=0, misses=9, maxsize=None, currsize=4)
V0801 01:28:54.231000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats get_implications: CacheInfo(hits=2, misses=3, maxsize=None, currsize=2)
V0801 01:28:54.231000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats get_axioms: CacheInfo(hits=3, misses=4, maxsize=None, currsize=1)
V0801 01:28:54.231000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats safe_expand: CacheInfo(hits=52, misses=28, maxsize=256, currsize=28)
V0801 01:28:54.231000 140445296705536 torch/fx/experimental/symbolic_shapes.py:116] lru_cache_stats uninteresting_files: CacheInfo(hits=39, misses=1, maxsize=None, currsize=1)
```

### Versions

Collecting environment information...
PyTorch version: 2.4.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.35

Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 2080
GPU 1: NVIDIA GeForce RTX 4080 SUPER

Nvidia driver version: 560.70
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             24
On-line CPU(s) list:                0-23
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen Threadripper 2920X 12-Core Processor
CPU family:                         23
Model:                              8
Thread(s) per core:                 2
Core(s) per socket:                 12
Socket(s):                          1
Stepping:                           2
BogoMIPS:                           7000.13
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr virt_ssbd arat npt nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload
Virtualization:                     AMD-V
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          384 KiB (12 instances)
L1i cache:                          768 KiB (12 instances)
L2 cache:                           6 MiB (12 instances)
L3 cache:                           8 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; untrained return thunk; SMT vulnerable
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.1
[pip3] pytorch-triton==2.1.0+bcad9dabe1
[pip3] torch==2.4.0+cu124
[pip3] torchviz==0.0.2
[pip3] triton==3.0.0
[conda] Could not collect

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99 @ezyang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4