Repository: pytorch
Issue: Graph breaks with HuggingFace Stable Diffusion
Labels: triaged,bug

### 🐛 Describe the bug

Using this as a tracking issue for graph breaks in HuggingFace Stable Diffusion. I'm going over the outputs from `explain` api, and trying to understand the different graph breaks and how to avoid them.

- [ ] Number of graphs do not match with number of graph break reasons. (Example script shows 105 graphs vs 75 reasons)
- [ ] call_function UserDefinedObjectVariable: source is a user defined class wrapping torch.Tensor as attribute. Can this be supported in dynamo?
    ```python
    @dataclass
    class Transformer2DModelOutput(BaseOutput):
        sample: torch.FloatTensor

    ...

    def some_func(...):
        ...
        return Transformer2DModelOutput(sample=output)
    ```
- [ ] data dependent operator: aten._local_scalar_dense.default: What's the recommended step to work around this issue? Can symbolic number help supporting this inside dynamo?
    ```
      File "/home/bowbao/diffusers/src/diffusers/schedulers/scheduling_pndm.py:337", line 337, in step_plms
        prev_sample = self._get_prev_sample(sample, timestep, prev_timestep, model_output)
      File "/home/bowbao/diffusers/src/diffusers/schedulers/scheduling_pndm.py:371", line 371, in _get_prev_sample
        alpha_prod_t = self.alphas_cumprod[timestep]
    ```
- [ ] Tensor.numpy: computation in numpy. 

### Error logs

```log
Dynamo produced 105 graphs with 104 graph break and 1367 ops
 Break reasons:

1. call_function UserDefinedObjectVariable(CLIPTokenizer) [ListVariable()] {'padding': ConstantVariable(str), 'max_length': ConstantVariable(int), 'truncation': ConstantVariable(bool), 'return_tensors': ConstantVariable(str)}
  File "/home/bowbao/diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py:285", line 285, in <graph break in _encode_prompt>
    return_tensors="pt",

2. call_function UserDefinedClassVariable() [] {'sample': TensorVariable()}
  File "/home/bowbao/diffusers/src/diffusers/models/attention.py:221", line 221, in forward
    return Transformer2DModelOutput(sample=output)

3. data dependent operator: aten._local_scalar_dense.default
  File "/home/bowbao/diffusers/src/diffusers/schedulers/scheduling_pndm.py:337", line 337, in step_plms
    prev_sample = self._get_prev_sample(sample, timestep, prev_timestep, model_output)
  File "/home/bowbao/diffusers/src/diffusers/schedulers/scheduling_pndm.py:371", line 371, in _get_prev_sample
    alpha_prod_t = self.alphas_cumprod[timestep]

4. data dependent operator: aten._local_scalar_dense.default
  File "/home/bowbao/diffusers/src/diffusers/schedulers/scheduling_pndm.py:372", line 372, in <graph break in _get_prev_sample>
    alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod

5. call_function UserDefinedClassVariable() [] {'sample': TensorVariable()}
  File "/home/bowbao/diffusers/src/diffusers/models/vae.py:586", line 586, in decode
    return DecoderOutput(sample=dec)

6. Tensor.numpy
  File "/home/bowbao/diffusers/src/diffusers/pipelines/stable_diffusion/safety_checker.py:56", line 56, in forward
    special_cos_dist = cosine_distance(image_embeds, self.special_care_embeds).cpu().float().numpy()

7. Tensor.numpy
  File "/home/bowbao/diffusers/src/diffusers/pipelines/stable_diffusion/safety_checker.py:57", line 57, in <graph break in forward>
    cos_dist = cosine_distance(image_embeds, self.concept_embeds).cpu().float().numpy()

TorchDynamo compilation metrics:
Function                                             Runtimes (s)
---------------------------------------------------  ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
convert_frame_assert.<locals>._convert_frame_assert  0.1300, 0.0086, 0.0030, 0.0030, 0.0000, 0.0042, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3416, 0.3279, 0.3279, 0.0081, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0115, 0.0000, 0.0120, 1.0618, 0.0001, 0.0001, 0.0000, 0.0000, 1.0440, 0.0074, 0.0036, 0.0057, 0.0000, 0.0029, 0.0016, 0.9361, 0.0241, 0.5783, 0.5774, 0.0163, 0.0016, 0.1725, 0.0350, 0.1384, 0.0017, 0.0000, 0.0000, 0.0000, 0.0026, 0.0029, 0.0000, 0.0352, 0.1380, 0.0060, 0.1713, 0.1389, 0.0014, 0.0021, 0.0023, 0.0351, 0.1375, 0.0060, 0.0377, 0.1374, 0.0014, 0.0021, 0.0023, 0.0347, 0.1368, 0.0060, 0.0704, 0.1692, 0.0024, 0.0012, 0.0007, 0.0349, 0.0031, 0.0023, 0.1366, 0.0014, 0.0021, 0.0023, 0.0350, 0.1263, 0.1744, 0.0375, 0.1372, 0.0384, 0.1364, 0.0381, 0.1366, 0.0082, 0.0375, 0.1365, 0.0375, 0.1366, 0.0374, 0.1367, 0.0082, 0.0375, 0.1374, 0.0374, 0.1375, 0.0375, 0.1362, 0.0016, 0.0000, 0.0021, 0.0023, 0.0123, 0.0158, 0.0062, 0.0074, 0.0065, 0.0153, 0.0034, 0.0017, 0.0000, 0.0022, 0.0022, 0.0187, 0.0165, 0.0034, 0.0178, 0.0155, 0.0034, 0.0209, 0.0189, 0.0036, 0.0227, 0.0213, 0.0035, 0.0238, 0.0217, 0.0036, 0.0237, 0.0216, 0.0034, 0.0240, 0.0216, 0.0035, 0.0235, 0.0219, 0.0034, 0.0236, 0.0217, 0.0036, 0.0235, 0.0216, 0.0035, 0.0237, 0.0216, 0.0035, 0.0233, 0.0216, 0.0036, 0.0240, 0.0216, 0.0034, 0.0233, 0.0217, 0.0035, 0.0236, 0.0217, 0.0034, 0.0235, 0.0215, 0.0035, 0.0239, 0.0218, 0.0036, 0.0235, 0.0222, 0.0035, 0.0236, 0.0216, 0.0037, 0.0236, 0.0217, 0.0034, 0.0238, 0.0215, 0.0035, 0.0235, 0.0216, 0.0036, 0.0240, 0.0216, 0.0034, 0.0234, 0.0216, 0.0035, 0.0245, 0.0218, 0.0035, 0.0235, 0.0216, 0.0037, 0.0240, 0.0216, 0.0035, 0.0236, 0.0215, 0.0036, 0.0238, 0.0218, 0.0035, 0.0241, 0.0216, 0.0035, 0.0237, 0.0215, 0.0037, 0.0237, 0.0217, 0.0035, 0.0236, 0.0216, 0.0036, 0.0243, 0.0221, 0.0036, 0.0236, 0.0215, 0.0037, 0.0238, 0.0220, 0.0037, 0.0237, 0.0215, 0.0037, 0.0234, 0.0219, 0.0035, 0.0237, 0.0216, 0.0035, 0.0235, 0.0219, 0.0035, 0.0236, 0.0216, 0.0035, 0.0237, 0.0216, 0.0035, 0.0237, 0.0215, 0.0035, 0.0236, 0.0216, 0.0035, 0.0238, 0.0216, 0.0035, 0.0236, 0.0215, 0.0035, 0.0238, 0.0218, 0.0035, 0.0234, 0.0215, 0.0036, 0.0237, 0.0221, 0.0035, 0.0235, 0.0216, 0.0150, 0.0035, 0.4317, 0.5953, 0.0018, 0.0000, 0.0022, 0.0024, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 2.0281, 0.0753, 0.0055, 0.0000, 0.0000, 0.0000
graphs: 105
reasons: 75
```

### Minified repro

```python
import torch
from diffusers import StableDiffusionPipeline
from torch import _dynamo as torchdynamo

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    revision="fp16",
)
pipe = pipe.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"

def fn(x):
    return pipe(x)

(
    explanation,
    out_guards,
    graphs,
    ops_per_graph,
    break_reasons,
    explanation_verbose,
) = torchdynamo.explain(fn, prompt)

print(explanation_verbose)  # Dynamo produced 105 graphs with 104 graph break and 1367 ops
print(f"graphs: {len(graphs)}")  # 105
print(f"reasons: {len(break_reasons)}")  # 75
```