Update utils/actor_utils.py

utils/actor_utils.py

@@ -1,79 +1,64 @@
-
 import torch
 
 
-def widen_tensor(datum, factor):
-
-    if len(datum.shape) == 0:
-        return datum
-
-    L = list(datum.shape)
-    a = [1, factor] + [1 for _ in range(len(L) - 1)]
+def widen_tensor(tensor, factor):
+    """ 
+    Duplicate a tensor `factor` times along the batch dimension.
+    """
+    if tensor.dim() == 0:
+        return tensor
 
-    datum = datum.unsqueeze(1)
-    datum = datum.repeat(*a)
+    shape = tensor.shape
+    repeat_dims = [1, factor] + [1] * (len(shape) - 1)
 
-    if len(L) > 1:
-        b = [L[0] * factor] + L[1:]
-    else:
-        b = [L[0] * factor]
-
-    datum = datum.reshape(*b)
-    return datum
+    expanded = tensor.unsqueeze(1).repeat(*repeat_dims)
+    new_shape = [shape[0] * factor] + list(shape[1:])
+    return expanded.view(*new_shape)
 
 
 def widen_data(actor, include_embeddings=True, include_projections=True):
+    """
+    Expand the actor's data `sample_size` times to support beam sampling.
+    """
 
-    F = dir(actor.fleet)
-    for s in F:
-        x = getattr(actor.fleet, s)
-        if isinstance(x, torch.Tensor):
-            if len(x.shape) > 0:
-                y = widen_tensor(x, factor=actor.sample_size)
-                setattr(actor.fleet, s, y)
+    def widen_attributes(obj):
+        for name, value in obj.__dict__.items():
+            if isinstance(value, torch.Tensor) and value.dim() > 0:
+                setattr(obj, name, widen_tensor(value, actor.sample_size))
 
-    G = dir(actor.graph)
-    for s in G:
-        x = getattr(actor.graph, s)
-        if isinstance(x, torch.Tensor):
-            if len(x.shape) > 0:
-                y = widen_tensor(x, factor=actor.sample_size)
-                setattr(actor.graph, s, y)
+    widen_attributes(actor.fleet)
+    widen_attributes(actor.graph)
 
-    actor.log_probs = widen_tensor(actor.log_probs, factor=actor.sample_size)
+    actor.log_probs = widen_tensor(actor.log_probs, actor.sample_size)
 
     if include_embeddings:
-        actor.node_embeddings = widen_tensor(actor.node_embeddings, factor=actor.sample_size)
+        actor.node_embeddings = widen_tensor(actor.node_embeddings, actor.sample_size)
 
     if include_projections:
-        def widen_projection(x, size):
-            if len(x.shape) > 3:
-                y = x.unsqueeze(2).repeat(1, 1, size, 1, 1)
-                return y.reshape(x.shape[0], x.shape[1] * size, x.shape[2], x.shape[3])
-            else:
-                return widen_tensor(x, size)
-
-        actor.node_projections = {key : widen_projection(actor.node_projections[key], actor.sample_size)
-                                 for key in actor.node_projections}
+        def widen_projection(x):
+            if x.dim() > 3:
+                # (heads, batch, graph, embed)
+                y = x.unsqueeze(2).repeat(1, 1, actor.sample_size, 1, 1)
+                return y.view(x.size(0), x.size(1) * actor.sample_size, x.size(2), x.size(3))
+            return widen_tensor(x, actor.sample_size)
 
+        actor.node_projections = {
+            key: widen_projection(value) for key, value in actor.node_projections.items()
+        }
 
 
 def select_data(actor, index, include_embeddings=True, include_projections=True):
-    m = index.max().item()
-
-    F = dir(actor.fleet)
-    for s in F:
-        x = getattr(actor.fleet, s)
-        if isinstance(x, torch.Tensor):
-            if (len(x.shape) > 0) and (x.shape[0] >= m):
-                setattr(actor.fleet, s, x[index])
-
-    G = dir(actor.graph)
-    for s in G:
-        x = getattr(actor.graph, s)
-        if isinstance(x, torch.Tensor):
-            if (len(x.shape) > 0) and (x.shape[0] >= m):
-                setattr(actor.graph, s, x[index])
+    """
+    Select a specific subset of data using `index` (e.g., for beam search pruning).
+    """
+
+    def select_attributes(obj):
+        for name, value in obj.__dict__.items():
+            if isinstance(value, torch.Tensor) and value.size(0) >= index.max().item():
+                setattr(obj, name, value[index])
+
+    select_attributes(actor.fleet)
+    select_attributes(actor.graph)
 
     actor.log_probs = actor.log_probs[index]
 
@@ -81,12 +66,11 @@ def select_data(actor, index, include_embeddings=True, include_projections=True)
         actor.node_embeddings = actor.node_embeddings[index]
 
     if include_projections:
-        def select_projection(x, index):
-            if len(x.shape) > 3:
-                return x[:,index,:,:]
-            else:
-                return x[index]
-
-        actor.node_projections = {key : select_projection(actor.node_projections[key], index)
-                                 for key in actor.node_projections}
-
+        def select_projection(x):
+            if x.dim() > 3:
+                return x[:, index, :, :]
+            return x[index]
+
+        actor.node_projections = {
+            key: select_projection(value) for key, value in actor.node_projections.items()
+        }