Fix GAT processors.

PiperOrigin-RevId: 427537373
google-deepmind · Feb 9, 2022 · d2731dd · d2731dd
1 parent 8287ae6
commit d2731dd
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/clrs/_src/processors.py b/clrs/_src/processors.py
@@ -75,8 +75,9 @@ def __call__(
     skip = hk.Linear(self.out_size)
 
     bias_mat = (adj - 1.0) * 1e9
-    bias_mat = jnp.tile(bias_mat, (1, 1, 1, self.nb_heads))  # [B, N, N, H]
-    bias_mat = jnp.transpose(bias_mat, (0, 3, 1, 2))         # [B, H, N, N]
+    bias_mat = jnp.tile(bias_mat[..., None],
+                        (1, 1, 1, self.nb_heads))     # [B, N, N, H]
+    bias_mat = jnp.transpose(bias_mat, (0, 3, 1, 2))  # [B, H, N, N]
 
     a_1 = hk.Linear(self.nb_heads)
     a_2 = hk.Linear(self.nb_heads)
@@ -169,8 +170,9 @@ def __call__(
     skip = hk.Linear(self.out_size)
 
     bias_mat = (adj - 1.0) * 1e9
-    bias_mat = jnp.tile(bias_mat, (1, 1, 1, self.nb_heads))  # [B, N, N, H]
-    bias_mat = jnp.transpose(bias_mat, (0, 3, 1, 2))         # [B, H, N, N]
+    bias_mat = jnp.tile(bias_mat[..., None],
+                        (1, 1, 1, self.nb_heads))     # [B, N, N, H]
+    bias_mat = jnp.transpose(bias_mat, (0, 3, 1, 2))  # [B, H, N, N]
 
     w_1 = hk.Linear(self.mid_size)
     w_2 = hk.Linear(self.mid_size)