Repository: pytorch
Issue: Issues with NJT constructed in graph w/ requires_grad=True and returned
Labels: triaged,module: nestedtensor,tensor subclass,oncall: pt2,dynamo-tensor-subclasses

### 🐛 Describe the bug

```python
import torch
from torch.nested._internal.nested_tensor import nested_view_from_values_offsets

def fn(values, offsets):
    return nested_view_from_values_offsets(values, offsets, max_seqlen=200, min_seqlen=1)

def get_values_offsets(batch_size=16, max_seqlen=200, inner_dim=64):
    lengths = torch.randint(1, max_seqlen, (batch_size,), dtype=torch.int32)
    offsets = torch.zeros((batch_size + 1,), dtype=torch.int32)
    torch.cumsum(lengths, dim=0, out=offsets[1:])
    values = torch.randn((offsets[-1].item(), inner_dim), dtype=torch.bfloat16, requires_grad=True)
    return values, offsets

values, offsets = get_values_offsets()
torch.compile(fn)(values, offsets)
```

First: getting this to run with requires_grad=False requires https://github.com/pytorch/pytorch/pull/124803 and https://github.com/pytorch/pytorch/pull/126563 (or other potential fixes - but these are confirmed to work with requires_grad=False)

After patching those, you'll get the following error (collapsed since the stack is pretty long):

<details>

```
Traceback (most recent call last):
  File "/home/dberard/local/scripts/nt_construct_requires_grad.py", line 15, in <module>
    torch.compile(fn)(values, offsets)
  File "/home/dberard/local/pytorch/torch/_dynamo/eval_frame.py", line 434, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 1121, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 948, in __call__
    result = self._inner_convert(
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
  File "/home/dberard/local/pytorch/torch/_utils_internal.py", line 85, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/home/dberard/local/pytorch/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/home/dberard/local/miniconda3/envs/pytorch/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/dberard/local/pytorch/torch/_dynamo/utils.py", line 233, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/dberard/local/pytorch/torch/_dynamo/bytecode_transformation.py", line 1270, in transform_code_object
    transformations(instructions, code_options)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2476, in run
    super().run()
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 904, in run
    while self.step():
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 816, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2667, in RETURN_VALUE
    self._return(inst)
  File "/home/dberard/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2652, in _return
    self.output.compile_subgraph(
  File "/home/dberard/local/pytorch/torch/_dynamo/output_graph.py", line 1102, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/home/dberard/local/miniconda3/envs/pytorch/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/dberard/local/pytorch/torch/_dynamo/output_graph.py", line 1324, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/dberard/local/pytorch/torch/_dynamo/utils.py", line 233, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/output_graph.py", line 1415, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/home/dberard/local/pytorch/torch/_dynamo/output_graph.py", line 1396, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/dberard/local/pytorch/torch/_dynamo/repro/after_dynamo.py", line 129, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/dberard/local/pytorch/torch/__init__.py", line 2143, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/home/dberard/local/miniconda3/envs/pytorch/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/dberard/local/pytorch/torch/_inductor/compile_fx.py", line 1550, in compile_fx
    return aot_autograd(
  File "/home/dberard/local/pytorch/torch/_dynamo/backends/common.py", line 69, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
  File "/home/dberard/local/pytorch/torch/_functorch/aot_autograd.py", line 990, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
  File "/home/dberard/local/pytorch/torch/_functorch/aot_autograd.py", line 979, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
  File "/home/dberard/local/pytorch/torch/_dynamo/utils.py", line 233, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_functorch/aot_autograd.py", line 708, in create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 266, in aot_dispatch_autograd
    fx_g, joint_inputs, maybe_subclass_meta = aot_dispatch_autograd_graph(
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 277, in aot_dispatch_autograd_graph
    fx_g = _create_graph(joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 46, in _create_graph
    fx_g = make_fx(
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 1428, in wrapped
    return make_fx_tracer.trace(f, *args)
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 1374, in trace
    return self._trace_inner(f, *args)
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 1361, in _trace_inner
    t = dispatch_trace(
  File "/home/dberard/local/pytorch/torch/_compile.py", line 31, in inner
    return disable_fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/_dynamo/eval_frame.py", line 601, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 642, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/dberard/local/pytorch/torch/_dynamo/eval_frame.py", line 601, in _fn
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/fx/_symbolic_trace.py", line 822, in trace
    (self.create_arg(fn(*args)),),
  File "/home/dberard/local/pytorch/torch/fx/_symbolic_trace.py", line 675, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 660, in wrapped
    out = f(*tensors)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 694, in joint_fn
    return inner_fn(flat_fn_maybe_joint, (primals, tangents), use_trace_joint=True)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 676, in inner_fn
    wrapped_outs = fn(*all_args)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 615, in joint_helper
    return _functionalized_f_helper(primals, tangents)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 389, in _functionalized_f_helper
    f_outs = fn(*f_args)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 257, in inner_fn_with_anomaly
    return inner_fn(*args)
  File "/home/dberard/local/pytorch/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 221, in inner_fn
    else out.view(tangent.shape)
  File "/home/dberard/local/pytorch/torch/fx/experimental/proxy_tensor.py", line 705, in __torch_function__
    return func(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/nested/_internal/nested_tensor.py", line 307, in __torch_function__
    return func(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/nested/_internal/nested_tensor.py", line 291, in __torch_dispatch__
    return fn(*args, **kwargs)
  File "/home/dberard/local/pytorch/torch/nested/_internal/ops.py", line 183, in inner
    return func(aten_op, *args, **kwargs)
  File "/home/dberard/local/pytorch/torch/nested/_internal/ops.py", line 908, in view_default
    raise RuntimeError(f"view(): cannot view shape {inp._size} as {size}")
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
RuntimeError: view(): cannot view shape (16, j4, 64) as [16, j2, 64]

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

</details>

The issue seems to be in AOTAutograd:
* collect_metadata_analysis.py first runs the function with some inputs to get the traced_tangents
* later (in the stack trace), we try to reshape the output to be the size expected for the tangents; but, since we ran through the function with different inputs (and a different offsets tensor), we get a different nested int in the shape.

### Versions

main branch

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer @YuqingJ @ezyang @albanD @anijain2305 @chauhang @penguinwu