jx_nest_base fails with TORCHDYNAMO_DYNAMIC_SHAPES=1: TypeError: reshape(): argument 'shape' must be tuple of SymInts, but found element of type float at pos 2 #93572
Description
🐛 Describe the bug
TORCHDYNAMO_DYNAMIC_SHAPES=1 time python benchmarks/dynamo/timm_models.py --accuracy --backend aot_eager --training --only jx_nest_base
fails with
File "/raid/ezyang/pytorch-scratch2/torch/_dynamo/utils.py", line 671, in wrap_fake_exception
return fn()
File "/raid/ezyang/pytorch-scratch2/torch/_dynamo/variables/tensor.py", line 132, in <lambda>
lambda: _run_node(tx.output, node, args, kwargs, nnmodule)
File "/raid/ezyang/pytorch-scratch2/torch/_dynamo/variables/tensor.py", line 53, in _run_node
return getattr(args[0], node.target)(*args[1:], **kwargs)
TypeError: reshape(): argument 'shape' must be tuple of SymInts, but found element of type float at pos 2
Error logs
https://gist.github.com/ezyang/f984fe6534db7b86fe2a24f7d99e8694
Minified repro
minifier did not work