Spaces:
Sleeping
Sleeping
| # Copyright 2022 DeepMind Technologies Limited. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Tests for compiler.craft_graph_to_model.""" | |
| from absl.testing import absltest | |
| from absl.testing import parameterized | |
| import networkx as nx | |
| from tracr.compiler import craft_graph_to_model | |
| from tracr.compiler import nodes | |
| from tracr.compiler import rasp_to_graph | |
| from tracr.craft import bases | |
| from tracr.craft.chamber import categorical_attn | |
| from tracr.craft.chamber import categorical_mlp | |
| from tracr.rasp import rasp | |
| class CraftAllocateModulesToLayersTest(parameterized.TestCase): | |
| def _get_dummy_block(self, block_type): | |
| if block_type == "ATTN": | |
| return categorical_attn.categorical_attn( | |
| query_space=bases.VectorSpaceWithBasis.from_names(["query"]), | |
| key_space=bases.VectorSpaceWithBasis.from_names(["bos", "key"]), | |
| value_space=bases.VectorSpaceWithBasis.from_names(["bos", "value"]), | |
| output_space=bases.VectorSpaceWithBasis.from_names(["output"]), | |
| bos_space=bases.VectorSpaceWithBasis.from_names(["bos"]), | |
| one_space=bases.VectorSpaceWithBasis.from_names(["one"]), | |
| attn_fn=lambda x, y: True, | |
| ) | |
| elif block_type == "MLP": | |
| return categorical_mlp.map_categorical_mlp( | |
| input_space=bases.VectorSpaceWithBasis.from_names(["input"]), | |
| output_space=bases.VectorSpaceWithBasis.from_names(["output"]), | |
| operation=lambda x: x, | |
| ) | |
| else: | |
| return None | |
| def test_get_longest_path_length_to_node_returns_expected_result(self): | |
| """Creates a graph and checks the longest path for each node.""" | |
| # Node IDs: | |
| # 0 -- 1 -- 2 -- 3 ------------ 4 | |
| # / / | |
| # 5 -- 6 ---------- 7 -- 8 -- 9 | |
| # | |
| # 10 | |
| # Expected return values: | |
| # 0 -- 1 -- 2 -- 3 ------------ 5 | |
| # / / | |
| # 0 -- 1 ---------- 2 -- 3 -- 4 | |
| # | |
| # -1 | |
| graph = nx.DiGraph() | |
| node_ids = list(range(11)) | |
| expected_results = [0, 1, 2, 3, 5, 0, 1, 2, 3, 4, -1] | |
| for node_id, res in zip(node_ids, expected_results): | |
| graph.add_node( | |
| node_id, **{ | |
| nodes.ID: node_id, | |
| nodes.EXPR: rasp.ConstantSOp(1), | |
| "expected_result": res | |
| }) | |
| graph.add_edge(0, 1) | |
| graph.add_edge(1, 2) | |
| graph.add_edge(2, 3) | |
| graph.add_edge(3, 4) | |
| graph.add_edge(5, 6) | |
| graph.add_edge(6, 7) | |
| graph.add_edge(7, 8) | |
| graph.add_edge(8, 9) | |
| graph.add_edge(6, 3) | |
| graph.add_edge(9, 4) | |
| sources = [graph.nodes[0], graph.nodes[5]] | |
| for node_id, node in graph.nodes.items(): | |
| result = craft_graph_to_model._get_longest_path_length_to_node( | |
| graph, sources, node) | |
| self.assertEqual(result, node["expected_result"]) | |
| def test_allocate_modules_to_layers_returns_expected_result(self): | |
| """Creates a graph and checks if the correct layer assignment is returned.""" | |
| # Computation Graph: | |
| # INPUT -- ATTN -- MLP -- ATTN ------ MLP -- OUTPUT | |
| # / / / | |
| # INPUT -- MLP --- MLP ATTN | |
| # \ / | |
| # ATTN | |
| # Node IDs: | |
| # 0 -- 1 -- 2 -- 3 -- 4 -- 5 | |
| # / / / | |
| # 6 -- 7 ---- 8 9 | |
| # \ / | |
| # 10 | |
| # Expected layer allocation: | |
| # -1 -- 0 -- 3 -- 4 -- 7 -- -1 | |
| # / / / | |
| # -1 -- 1 --- 3 6 | |
| # \ / | |
| # 4 | |
| graph = nx.DiGraph() | |
| node_ids = list(range(11)) | |
| types = [ | |
| "INPUT", "ATTN", "MLP", "ATTN", "MLP", "OUTPUT", "INPUT", "MLP", "MLP", | |
| "ATTN", "ATTN" | |
| ] | |
| expected_results = [-1, 0, 3, 4, 7, -1, -1, 1, 3, 6, 4] | |
| for node_id, node_type, res in zip(node_ids, types, expected_results): | |
| graph.add_node( | |
| node_id, **{ | |
| nodes.ID: node_id, | |
| nodes.EXPR: rasp.ConstantSOp(1), | |
| nodes.MODEL_BLOCK: self._get_dummy_block(node_type), | |
| "expected_result": res | |
| }) | |
| graph.add_edge(0, 1) | |
| graph.add_edge(1, 2) | |
| graph.add_edge(2, 3) | |
| graph.add_edge(3, 4) | |
| graph.add_edge(4, 5) | |
| graph.add_edge(6, 7) | |
| graph.add_edge(7, 2) | |
| graph.add_edge(7, 8) | |
| graph.add_edge(8, 3) | |
| graph.add_edge(8, 10) | |
| graph.add_edge(9, 4) | |
| graph.add_edge(10, 9) | |
| craft_graph = rasp_to_graph.ExtractRaspGraphOutput( | |
| graph=graph, | |
| sink=graph.nodes[10], | |
| sources=[graph.nodes[0], graph.nodes[6]]) | |
| layer_allocation = craft_graph_to_model._allocate_modules_to_layers( | |
| craft_graph.graph, craft_graph.sources) | |
| for node_id, node in graph.nodes.items(): | |
| self.assertEqual(layer_allocation[node_id], node["expected_result"]) | |
| def test_allocate_modules_to_layers_returns_expected_result_for_chain(self): | |
| """Tests a chain of alternating attention layers and MLPs.""" | |
| # Computation Graph: | |
| # INPUT -- ATTN -- MLP -- ATTN -- MLP -- OUTPUT | |
| # Node IDs: | |
| # 0 -- 1 -- 2 -- 3 -- 4 -- 5 | |
| # Expected layer allocation: | |
| # -1 -- 0 -- 1 -- 2 -- 3 -- -1 | |
| graph = nx.DiGraph() | |
| node_ids = list(range(11)) | |
| types = ["INPUT", "ATTN", "MLP", "ATTN", "MLP", "OUTPUT"] | |
| expected_results = [-1, 0, 1, 2, 3, -1] | |
| for node_id, node_type, res in zip(node_ids, types, expected_results): | |
| graph.add_node( | |
| node_id, **{ | |
| nodes.ID: node_id, | |
| nodes.EXPR: rasp.ConstantSOp(1), | |
| nodes.MODEL_BLOCK: self._get_dummy_block(node_type), | |
| "expected_result": res | |
| }) | |
| graph.add_edge(0, 1) | |
| graph.add_edge(1, 2) | |
| graph.add_edge(2, 3) | |
| graph.add_edge(3, 4) | |
| graph.add_edge(4, 5) | |
| craft_graph = rasp_to_graph.ExtractRaspGraphOutput( | |
| graph=graph, sink=graph.nodes[5], sources=[graph.nodes[0]]) | |
| layer_allocation = craft_graph_to_model._allocate_modules_to_layers( | |
| craft_graph.graph, craft_graph.sources) | |
| for node_id, node in graph.nodes.items(): | |
| self.assertEqual(layer_allocation[node_id], node["expected_result"]) | |
| if __name__ == "__main__": | |
| absltest.main() | |