AllLatticeEdges.get_edges written, not tested

understanding-search · May 25, 2024 · ba6404a · ba6404a
1 parent f680b1f
commit ba6404a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/maze_dataset/tokenization/maze_tokenizer.py b/maze_dataset/tokenization/maze_tokenizer.py
@@ -749,7 +749,8 @@ class AllLatticeEdges(EdgeSubset):
         def get_edges(
             maze: LatticeMaze, 
         ) -> ConnectionArray:
-            return lattice_connection_array(maze.grid_n)
+            lattice_edges: ConnectionArray = lattice_connection_array(maze.grid_n)
+            return np.append(lattice_edges, np.flip(lattice_edges, axis=1), axis=0)
 
 
     @serializable_dataclass(frozen=True, kw_only=True)