[aseemw]

This PR fixes a failure in graph-mode quantization when a model contains the same op written in two different ways, for example one add written as x + y and another as torch.add(x, y). Quantizer.prepare() currently fails on such models with:

RuntimeError: prepare_qat_pt2e call failed, with error: Expected exactly 1 call function node in source partition but got [add, add_1].

A minimal example that reproduces it:

class Model(torch.nn.Module):
    def forward(self, x1, x2):
        out = x1 + x2
        out = torch.add(out, x2)
        return out

Cause

Op-list patterns (add, mul, cat, pooling, ...) were matched with torchao's find_sequential_partitions, which gets the ops from torch's get_source_partitions. That groups nodes into partitions based on the first element of node.meta["torch_fn"], a name made of the called function's __name__ and a count of that function's calls. x + y (traced as torch.Tensor.add) and torch.add are different functions with the same __name__, counted separately, so both nodes are named "add_1" and end up in a single partition. The annotation expects one op per partition and raises.

Fix

• Removes the dependency on torchao.quantization.pt2e.find_sequential_partitions and adds a local _find_sequential_partitions. It matches ops node by node, using each node's op type, and returns results in the same format, but with every op in its own SourcePartition, so both adds in the example above are matched and quantized.
• Removes the special case in _get_call_function_node_from_partition for partitions of SymInt shape-assert nodes, which can no longer occur.
• Corrects the tuple[SourcePartition] type annotations to tuple[SourcePartition, ...], since a match holds one partition per op in the pattern.

[NehalBhandari]

Looks good with one minor comment