Update to PyTorch 2.2

[awf]

A number of changes needed to work with PT2.2+. Primarily, there's now (since commit) no good way to intercept module calls, so we instead replace all nn.Modules with trivial subclasses, making them "user" modules.

Passes tests on asses tests on 2.1,2.2 (stable),2.4 (nightly) and examples/scale_analysis.py produces identical output (with torch.manual_seed).

This implementation is somewhat faster, as it does less work in the patched forward functions.

One source of changes is that the node naming post 2.2 better reflects the input code.
For example, input code

def forward(self, idxs: Tensor) -> Tuple[Tensor, Tensor]:  # pragma: no cover
    # idxs has 0 args -> shouldn't be pruned
    x = self.emb(idxs)  # emb has 1 float arg (weights) -> depends on tol
    _x = x.flatten(start_dim=0, end_dim=-1)  # 1 float, same scale -> prune
    x = _x.view(x.shape)  # 1 float arg, same scale -> prune
    y = self.linear(x)  # scale changes -> shouldn't be pruned
    scores = F.softmax(y, dim=-1)  # scale changes -> shouldn't be pruned
    top_idx = torch.argmax(scores, dim=-1)  # not float -> shouldn't be pruned
    top_idx = torch.unsqueeze(top_idx, -1)  # not float -> shouldn't be pruned
    top_score_x = torch.gather(x, -1, top_idx)  # small change -> depends on tol
    top_score_x += randn_like(top_score_x)  # 2 floats, same scale -> no prune
    return top_score_x, top_idx

Became, pre 2.2, where the variable names are derived from the operator:

# Example of a pre-2.2 captured graph
def forward(self, L_idxs_ : torch.Tensor):
    l_idxs_ = L_idxs_
    l__self___emb_weight = foo.L__self___emb_weight
    embedding = torch.nn.functional.embedding(l_idxs_, l__self___emb_weight, None, None, 2.0, False, False);  l_idxs_ = l__self___emb_weight = None
    flatten = embedding.flatten(start_dim = 0, end_dim = -1);  embedding = None
    view = flatten.view((8, 32, 64));  flatten = None
    l__self___linear_weight = foo.L__self___linear_weight
    l__self___linear_bias = foo.L__self___linear_bias
    linear = torch._C._nn.linear(view, l__self___linear_weight, l__self___linear_bias);  l__self___linear_weight = l__self___linear_bias = None
    softmax = torch.nn.functional.softmax(linear, dim = -1);  linear = None
    argmax = torch.argmax(softmax, dim = -1);  softmax = None
    unsqueeze = torch.unsqueeze(argmax, -1);  argmax = None
    gather = torch.gather(view, -1, unsqueeze);  view = None
    randn_like = torch.randn_like(gather)
    gather += randn_like;  iadd = gather;  gather = randn_like = None
    return (iadd, unsqueeze)

and are now

# Example of a post-2.2 captured graph
def forward(self, L_idxs_ : torch.Tensor):
    l_idxs_ = L_idxs_
    l__self___emb_weight = foo.L__self___emb_weight
    x = torch.nn.functional.embedding(l_idxs_, l__self___emb_weight, None, None, 2.0, False, False);  l_idxs_ = l__self___emb_weight = None
    _x = x.flatten(start_dim = 0, end_dim = -1);  x = None
    x_1 = _x.view((8, 32, 64));  _x = None
    l__self___linear_weight = foo.L__self___linear_weight
    l__self___linear_bias = foo.L__self___linear_bias
    y = torch._C._nn.linear(x_1, l__self___linear_weight, l__self___linear_bias);  l__self___linear_weight = l__self___linear_bias = None
    scores = torch.nn.functional.softmax(y, dim = -1);  y = None
    top_idx = torch.argmax(scores, dim = -1);  scores = None
    top_idx_1 = torch.unsqueeze(top_idx, -1);  top_idx = None
    top_score_x = torch.gather(x_1, -1, top_idx_1);  x_1 = None
    randn_like = torch.randn_like(top_score_x)
    top_score_x += randn_like;  top_score_x_1 = top_score_x;  top_score_x = randn_like = None
    return (top_score_x_1, top_idx_1)

[awf]

This is suppressing a warning about setting yticklabels without setting ticks

[thecharlieblake]

Thanks for looking into this Andrew, pleased with the approach we've converged to (and learned a few things about deepcopy and pickling in the process).

I'll leave merging to you in case there are any other adjustments