finalize branch

tesseract/client/moe.py

@@ -147,7 +147,7 @@ class RemoteMixtureOfExperts(nn.Module):
 
     def compute_expert_scores(
             self, grid_scores: List[torch.Tensor], batch_experts: List[List[RemoteExpert]]) -> torch.Tensor:
-        """ TODO docstring here """
+        """ TODO(jheuristic) docstring here """
         expert_counts = list(map(len, batch_experts))
         batch_size = len(batch_experts)
         max_num_experts = max(expert_counts)

tesseract/utils/threading.py

@@ -64,4 +64,4 @@ def run_and_await_k(jobs: List[callable], k: int,
         for future, index in future_to_ix.items():
             future.cancel()
             outputs[index] = future.result() if not future.exception() else future.exception()
-    return outputs
+    return outputs