chore: import upstream snapshot with attribution

examples/pytorch/diffpool/model/tensorized_layers/assignment.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn as nn
+from torch.autograd import Variable
+from torch.nn import functional as F
+
+from model.tensorized_layers.graphsage import BatchedGraphSAGE
+
+
+class DiffPoolAssignment(nn.Module):
+    def __init__(self, nfeat, nnext):
+        super().__init__()
+        self.assign_mat = BatchedGraphSAGE(nfeat, nnext, use_bn=True)
+
+    def forward(self, x, adj, log=False):
+        s_l_init = self.assign_mat(x, adj)
+        s_l = F.softmax(s_l_init, dim=-1)
+        return s_l