Update graph_decoder/transformer.py

graph_decoder/transformer.py

@@ -2,39 +2,7 @@ import torch
 import torch.nn as nn
 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
+from .diffusion_utils import PlaceHolder
 
 def modulate(x, shift, scale):
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)