Repro: control_deps mixed-validity partitioner crash

repro_wait_stream.py

"""
User-code repro: control_deps mixed-validity partitioner crash.

Crashes WITHOUT the fix in partitioners.py with:
    AssertionError: Node _unsafe_view was invalid, but is output

Root cause:
  1. s1.wait_stream(default_stream) generates a wait_stream op in the graph.
  2. _collect_wait_stream_forward_deps scans ALL subsequent ops on stream s1 --
     including backward ops -- and collects their inputs defined before the
     wait_stream.  Tangent placeholders (defined at graph top, used by backward
     ops on s1) get collected as extra deps.
  3. _wrap_sync_node threads these as pass-throughs alongside forward primal
     weights.  The resulting control_deps has mixed forward+backward deps.
  4. Forward ops on s1 (linear1's mm) read their weight through a getitem from
     this control_deps.
  5. During forward-graph extraction: tangent placeholder -> InvalidNode ->
     control_deps invalid -> getitem invalid -> mm invalid -> but mm's output
     (_unsafe_view) IS a forward output -> crash.

This is the same mechanism as MAST job aps-dli7-4b52ce6593 (APS Ads ATLAS),
which crashed with "Node addmm_3 was invalid, but is output".

NOTE: With the fix applied, the partitioning succeeds but a separate runtime
error ("User object is no longer alive") occurs in the stream registry. That is
a different bug (addressed by D108769923 Bug 2), not the partitioner fix.
"""
import torch
import torch.nn as nn


class WaitStreamModel(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.linear1 = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        s1 = torch.cuda.Stream()
        default_stream = torch.cuda.current_stream()
        s1.wait_stream(default_stream)
        with torch.cuda.stream(s1):
            h = self.linear1(x)
        event = torch.cuda.Event()
        event.record(s1)
        event.wait()
        return h


if not torch.cuda.is_available():
    print("No CUDA")
    exit()

model = WaitStreamModel(64).cuda()
x = torch.randn(2, 16, 64, device="cuda", requires_grad=True)

# Eager sanity
out = model(x)
out.sum().backward()
print(f"Eager OK: grad_norm={x.grad.norm().item():.4f}")

for p in model.parameters():
    p.grad = None
x.grad = None

# Compiled: triggers the partitioner crash without the fix
torch._dynamo.reset()
compiled = torch.compile(model, backend="aot_eager")
try:
    out = compiled(x)
    out.sum().backward()
    print(f"Compiled OK: grad_norm={x.grad.norm().item():.4f}")
except Exception as e:
    print(f"BUG: {type(e).__name__}: {e}")