Update graph_decoder/transformer.py

graph_decoder/transformer.py

@@ -4,6 +4,38 @@ from .layers import Attention, MLP
 from .conditions import TimestepEmbedder, ConditionEmbedder
 # from .diffusion_utils import PlaceHolder
 
+#### graph utils
+class PlaceHolder:
+    def __init__(self, X, E, y):
+        self.X = X
+        self.E = E
+        self.y = y
+
+    def type_as(self, x: torch.Tensor, categorical: bool = False):
+        """Changes the device and dtype of X, E, y."""
+        self.X = self.X.type_as(x)
+        self.E = self.E.type_as(x)
+        if categorical:
+            self.y = self.y.type_as(x)
+        return self
+
+    def mask(self, node_mask, collapse=False):
+        x_mask = node_mask.unsqueeze(-1)  # bs, n, 1
+        e_mask1 = x_mask.unsqueeze(2)  # bs, n, 1, 1
+        e_mask2 = x_mask.unsqueeze(1)  # bs, 1, n, 1
+
+        if collapse:
+            self.X = torch.argmax(self.X, dim=-1)
+            self.E = torch.argmax(self.E, dim=-1)
+
+            self.X[node_mask == 0] = -1
+            self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = -1
+        else:
+            self.X = self.X * x_mask
+            self.E = self.E * e_mask1 * e_mask2
+            assert torch.allclose(self.E, torch.transpose(self.E, 1, 2))
+        return self
+
 def modulate(x, shift, scale):
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
 
@@ -98,8 +130,7 @@ class Transformer(nn.Module):
 
         # X: B * N * dx, E: B * N * N * de
         X, E = self.output_layer(X, X_in, E_in, c, t, node_mask)
-        # return PlaceHolder(X=X, E=E, y=None).mask(node_mask)
-        return X, E
+        return PlaceHolder(X=X, E=E, y=None).mask(node_mask)
 
 class Block(nn.Module):
     def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):