diff --git a/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py b/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py index 49e4c76cf..406e9660e 100644 --- a/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py +++ b/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py @@ -51,7 +51,6 @@ class BitwidthMode(Enum): single-precision nodes. To compute custom single precision configuration, use QCustom. """ Float = auto() - Q8Bit = auto() QMaxBit = auto() QMinBit = auto() QCustom = auto() @@ -573,7 +572,7 @@ def compute_node_bops(self, not (a_node.is_activation_quantization_enabled() or n.is_weights_quantization_enabled(kernel_attr))): return 0 - act_qc = act_qcs.get(a_node.name) if act_qcs else None + act_qc = self._extract_qc(a_node, act_qcs) a_nbits = self._get_activation_nbits(a_node, bitwidth_mode, act_qc) w_nbits = self._get_weight_nbits(n, kernel_attr, bitwidth_mode, w_qc) node_bops = a_nbits * w_nbits * node_mac @@ -708,23 +707,20 @@ def _get_activation_nbits(self, Returns: Activation bit-width. """ + n = self.graph.retrieve_preserved_quantization_node(n) if act_qc: assert bitwidth_mode == BitwidthMode.QCustom return act_qc.activation_n_bits if act_qc.quant_mode == ActivationQuantizationMode.QUANT else FLOAT_BITWIDTH - if bitwidth_mode == BitwidthMode.Float or not (n.is_activation_quantization_enabled() or - n.is_quantization_preserving()): + if bitwidth_mode == BitwidthMode.Float or not n.is_activation_quantization_enabled(): return FLOAT_BITWIDTH - if bitwidth_mode == BitwidthMode.Q8Bit: - return 8 - if bitwidth_mode in self._bitwidth_mode_fn: candidates_nbits = [c.activation_quantization_cfg.activation_n_bits for c in n.candidates_quantization_cfg] return self._bitwidth_mode_fn[bitwidth_mode](candidates_nbits) if bitwidth_mode in [BitwidthMode.QCustom, BitwidthMode.QDefaultSP]: - qcs = self.graph.retrieve_preserved_quantization_node(n).get_unique_activation_candidates() + qcs = n.get_unique_activation_candidates() if len(qcs) != 1: raise ValueError(f'Could not retrieve the activation quantization candidate for node {n} ' f'as it has {len(qcs)}!=1 unique candidates.') @@ -760,9 +756,6 @@ def _get_weight_nbits(cls, if bitwidth_mode == BitwidthMode.Float or not n.is_weights_quantization_enabled(w_attr): return FLOAT_BITWIDTH - if bitwidth_mode == BitwidthMode.Q8Bit: - return 8 - node_qcs = n.get_unique_weights_candidates(w_attr) w_qcs = [qc.weights_quantization_cfg.get_attr_config(w_attr) for qc in node_qcs] if bitwidth_mode in cls._bitwidth_mode_fn: diff --git a/tests_pytest/_test_util/graph_builder_utils.py b/tests_pytest/_test_util/graph_builder_utils.py index 7d07f0898..3619898ad 100644 --- a/tests_pytest/_test_util/graph_builder_utils.py +++ b/tests_pytest/_test_util/graph_builder_utils.py @@ -83,6 +83,7 @@ def build_nbits_qc(a_nbits=8, a_enable=True, w_attr=None, pos_attr=(32, False, ( Final name can be passed along with convert_canonical_attr=False. pos_attr: quantization configuration for positional weights in format (nbits, q enabled, indices). convert_canonical_attr: whether to convert w_attr keys to full names. + q_preserving: Whether node is quantization preserving. Returns: diff --git a/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py b/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py index 906fc8153..2ee1559ba 100644 --- a/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py +++ b/tests_pytest/common_tests/unit_tests/core/mixed_precision/resource_utilization_tools/test_resource_utilization_calculator.py @@ -37,7 +37,12 @@ BM = BitwidthMode TIC = TargetInclusionCriterion -_identity_func = lambda x: x + +@pytest.fixture +def graph_mock(): + """ Basic Graph mock with basic retrieve_preserved_quantization_node operation for handling non + quantization preserving nodes. """ + return Mock(spec_set=Graph, nodes=[], retrieve_preserved_quantization_node=lambda x: x) class TestUtilization: @@ -112,45 +117,45 @@ def test_compute_ru_all_targets(self, detailed): @pytest.mark.parametrize('detailed', [True, False]) def test_compute_ru_w(self, detailed): ru_calc = self.ru_calc - ret = ru_calc.compute_resource_utilization(TIC.Any, BM.Q8Bit, ru_targets=[RUTarget.WEIGHTS], + ret = ru_calc.compute_resource_utilization(TIC.Any, BM.QDefaultSP, ru_targets=[RUTarget.WEIGHTS], return_detailed=detailed) - ru_calc.compute_weights_utilization.assert_called_once_with(TIC.Any, BM.Q8Bit, None) + ru_calc.compute_weights_utilization.assert_called_once_with(TIC.Any, BM.QDefaultSP, None) ru_calc.compute_activations_utilization.assert_not_called() ru_calc.compute_bops.assert_not_called() - self._validate(ret, detailed, ResourceUtilization(weights_memory=42)) + self._validate(ret, detailed, ResourceUtilization(weights_memory=21)) @pytest.mark.parametrize('detailed', [True, False]) def test_compute_ru_act(self, detailed): ru_calc = self.ru_calc - ret = self.ru_calc.compute_resource_utilization(TIC.Any, BM.Q8Bit, ru_targets=[RUTarget.ACTIVATION], + ret = self.ru_calc.compute_resource_utilization(TIC.Any, BM.QDefaultSP, ru_targets=[RUTarget.ACTIVATION], return_detailed=detailed) - ru_calc.compute_activations_utilization.assert_called_once_with(TIC.Any, BM.Q8Bit, None) + ru_calc.compute_activations_utilization.assert_called_once_with(TIC.Any, BM.QDefaultSP, None) ru_calc.compute_weights_utilization.assert_not_called() ru_calc.compute_bops.assert_not_called() - self._validate(ret, detailed, ResourceUtilization(activation_memory=742)) + self._validate(ret, detailed, ResourceUtilization(activation_memory=671)) @pytest.mark.parametrize('detailed', [True, False]) def test_compute_ru_total(self, detailed): ru_calc = self.ru_calc - ret = ru_calc.compute_resource_utilization(TIC.Any, BM.Q8Bit, ru_targets=[RUTarget.TOTAL], + ret = ru_calc.compute_resource_utilization(TIC.Any, BM.QDefaultSP, ru_targets=[RUTarget.TOTAL], return_detailed=detailed) - ru_calc.compute_activations_utilization.assert_called_once_with(TIC.Any, BM.Q8Bit, None) - ru_calc.compute_weights_utilization.assert_called_once_with(TIC.Any, BM.Q8Bit, None) + ru_calc.compute_activations_utilization.assert_called_once_with(TIC.Any, BM.QDefaultSP, None) + ru_calc.compute_weights_utilization.assert_called_once_with(TIC.Any, BM.QDefaultSP, None) ru_calc.compute_bops.assert_not_called() - self._validate(ret, detailed, ResourceUtilization(total_memory=742+42)) + self._validate(ret, detailed, ResourceUtilization(total_memory=671+21)) @pytest.mark.parametrize('detailed', [True, False]) def test_compute_ru_bops(self, detailed): ru_calc = self.ru_calc - ret = ru_calc.compute_resource_utilization(TIC.AnyQuantized, BM.Q8Bit, ru_targets=[RUTarget.BOPS], + ret = ru_calc.compute_resource_utilization(TIC.AnyQuantized, BM.QDefaultSP, ru_targets=[RUTarget.BOPS], return_detailed=detailed) - ru_calc.compute_bops.assert_called_once_with(TIC.AnyQuantized, BM.Q8Bit, act_qcs=None, w_qcs=None) + ru_calc.compute_bops.assert_called_once_with(TIC.AnyQuantized, BM.QDefaultSP, act_qcs=None, w_qcs=None) ru_calc.compute_activations_utilization.assert_not_called() ru_calc.compute_weights_utilization.assert_not_called() - self._validate(ret, detailed, ResourceUtilization(bops=42*8*8)) + self._validate(ret, detailed, ResourceUtilization(bops=42*8*4)) def test_compute_ru_custom_w_qcs(self): ru_calc = self.ru_calc @@ -224,7 +229,17 @@ def test_get_a_nbits_configurable(self, graph_mock, fw_impl_mock, fw_info_mock): assert ru_calc._get_activation_nbits(node, BM.Float, None) == FLOAT_BITWIDTH assert ru_calc._get_activation_nbits(node, BM.QMinBit, None) == 2 assert ru_calc._get_activation_nbits(node, BM.QMaxBit, None) == 7 - assert ru_calc._get_activation_nbits(node, BM.Q8Bit, None) == 8 + + def test_get_a_nbits_configurable_quantization_preserving(self, graph_mock, fw_impl_mock, fw_info_mock): + ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock) + node = build_node(qcs=[build_qc(7, a_enable=False, q_preserving=True), + build_qc(4, a_enable=False, q_preserving=True), + build_qc(2, a_enable=False, q_preserving=True)]) + anode = build_node(qcs=[build_qc(17), build_qc(4), build_qc(1)]) + graph_mock.retrieve_preserved_quantization_node = lambda x: anode + assert ru_calc._get_activation_nbits(node, BM.Float, None) == FLOAT_BITWIDTH + assert ru_calc._get_activation_nbits(node, BM.QMinBit, None) == 1 + assert ru_calc._get_activation_nbits(node, BM.QMaxBit, None) == 17 @pytest.mark.parametrize('node', [ build_node(qcs=[build_qc(42)]), @@ -232,10 +247,9 @@ def test_get_a_nbits_configurable(self, graph_mock, fw_impl_mock, fw_info_mock): ]) def test_get_a_nbits_nonconfigurable(self, graph_mock, fw_impl_mock, fw_info_mock, node): ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock) - for bm in set(BitwidthMode) - {BM.Float, BM.Q8Bit}: + for bm in set(BitwidthMode) - {BM.Float}: assert ru_calc._get_activation_nbits(node, bm, None) == 42 assert ru_calc._get_activation_nbits(node, BM.Float, None) == FLOAT_BITWIDTH - assert ru_calc._get_activation_nbits(node, BM.Q8Bit, None) == 8 @pytest.mark.parametrize('node, qc, exp_nbit', [ (build_node(qcs=[build_qc(4)]), build_qc(17), 17), @@ -246,6 +260,16 @@ def test_get_a_nbits_custom(self, graph_mock, fw_impl_mock, fw_info_mock, node, ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock) assert ru_calc._get_activation_nbits(node, BM.QCustom, qc.activation_quantization_cfg) == exp_nbit + @pytest.mark.parametrize('anode, node, qc, exp_nbit', [ + (None, build_node(qcs=[build_qc(4)]), build_qc(17, False, q_preserving=True), 32), + (build_node(qcs=[build_qc(3)]), build_node(qcs=[build_qc(4, False, q_preserving=True)]), None, 3) + ]) + def test_get_a_nbits_custom_quantization_preserving(self, graph_mock, fw_impl_mock, fw_info_mock, anode, node, qc, exp_nbit): + graph_mock.retrieve_preserved_quantization_node = lambda x: anode + ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock) + acs = None if qc is None else qc.activation_quantization_cfg + assert ru_calc._get_activation_nbits(node, BM.QCustom, acs) == exp_nbit + @pytest.mark.parametrize('bm', list(BM)) def test_get_a_nbits_non_q(self, graph_mock, fw_impl_mock, fw_info_mock, bm): node = build_node(qcs=[build_qc(a_enable=False)]) @@ -270,8 +294,9 @@ def test_get_target_activation_nodes(self, graph_mock, fw_impl_mock, fw_info_moc mp = build_node('n4', qcs=[build_qc(4), build_qc(2)], reuse=True) noq = build_node('noq', qcs=[build_qc(4, False, w_attr={'foo': (8, True)}), build_qc(4, False, w_attr={'foo': (4, True)})]) + qp = build_node('qp', qcs=[build_qc(4, False, q_preserving=True)]) - graph_mock.nodes = [sp1, sp2, sp3, mp, noq] + graph_mock.nodes = [sp1, sp2, sp3, mp, noq, qp] ru_calc = ResourceUtilizationCalculator(graph_mock, fw_impl_mock, fw_info_mock) assert len(TIC) == 4, 'enum changed, update tests' @@ -281,11 +306,11 @@ def test_get_target_activation_nodes(self, graph_mock, fw_impl_mock, fw_info_moc assert ru_calc._get_target_activation_nodes(TIC.QNonConfigurable, include_reused=True) == [sp2, sp3] assert ru_calc._get_target_activation_nodes(TIC.QNonConfigurable, include_reused=False) == [sp2] - assert ru_calc._get_target_activation_nodes(TIC.AnyQuantized, include_reused=True) == [sp1, sp2, sp3, mp] - assert ru_calc._get_target_activation_nodes(TIC.AnyQuantized, include_reused=False) == [sp1, sp2] + assert ru_calc._get_target_activation_nodes(TIC.AnyQuantized, include_reused=True) == [sp1, sp2, sp3, mp, qp] + assert ru_calc._get_target_activation_nodes(TIC.AnyQuantized, include_reused=False) == [sp1, sp2, qp] - assert ru_calc._get_target_activation_nodes(TIC.Any, include_reused=True) == [sp1, sp2, sp3, mp, noq] - assert ru_calc._get_target_activation_nodes(TIC.Any, include_reused=False) == [sp1, sp2, noq] + assert ru_calc._get_target_activation_nodes(TIC.Any, include_reused=True) == [sp1, sp2, sp3, mp, noq, qp] + assert ru_calc._get_target_activation_nodes(TIC.Any, include_reused=False) == [sp1, sp2, noq, qp] # explicit nodes list assert ru_calc._get_target_activation_nodes(TIC.QNonConfigurable, include_reused=True, nodes=[sp1, sp2, sp3]) == [sp2, sp3] @@ -320,9 +345,6 @@ def test_compute_node_activation_tensor_utilization(self, graph_mock, fw_impl_mo # not a target node res = ru_calc.compute_node_activation_tensor_utilization(noq, TIC.AnyQuantized, BM.QCustom, custom_qc) assert res == Utilization(0, 0) - # no target - res = ru_calc.compute_node_activation_tensor_utilization(noq, None, BM.Q8Bit) - assert res == Utilization(135, 540.) @pytest.mark.parametrize('bitmode', set(BM)-{BM.QCustom}) def test_compute_node_activation_tensor_utilization_errors(self, graph_mock, fw_impl_mock, fw_info_mock, bitmode): @@ -467,11 +489,14 @@ def prepare_compute_cuts(self, graph_mock, fw_impl_mock, fw_info_mock, mocker): noq = build_node('noq', qcs=[build_qc(6, False)], output_shape=(None, 300)) sp = build_node('sp', qcs=[build_qc(3)], output_shape=(None, 20, 10)) mp2 = build_node('mp2', qcs=[build_qc(2), build_qc(4)], output_shape=(None, 150)) + qp = build_node('qp', qcs=[build_qc(2, a_enable=False, q_preserving=True), + build_qc(4, a_enable=False, q_preserving=True)], output_shape=(None, 150)) - nodes = [mp_reuse, mp, noq, sp, mp2] + nodes = [mp_reuse, mp, noq, sp, mp2, qp] graph_mock.nodes = nodes # use the Graph original method (need to bind it to graph_mock instance) graph_mock.find_node_by_name = MethodType(Graph.find_node_by_name, graph_mock) + graph_mock.retrieve_preserved_quantization_node = lambda x: mp2 if x.name == 'qp' else x # we should not use total size, setting it to bad number cut_elems1 = MemoryElements(elements={ActivationMemoryTensor(mp_reuse.output_shape, 'mp_reuse', 0)}, total_size=-1) @@ -481,7 +506,8 @@ def prepare_compute_cuts(self, graph_mock, fw_impl_mock, fw_info_mock, mocker): ActivationMemoryTensor(sp.output_shape, 'sp', 0)}, total_size=-1) cut_elems3 = MemoryElements(elements={ActivationMemoryTensor(sp.output_shape, 'sp', 0), ActivationMemoryTensor(noq.output_shape, 'noq', 0)}, total_size=-1) - cut_elems4 = MemoryElements(elements={ActivationMemoryTensor(mp2.output_shape, 'mp2', 0)}, total_size=-1) + cut_elems4 = MemoryElements(elements={ActivationMemoryTensor(mp2.output_shape, 'mp2', 0), + ActivationMemoryTensor(mp2.output_shape, 'qp', 0)}, total_size=-1) cuts = [Cut([], set(), mem_elements=cut_elems) for cut_elems in [cut_elems1, cut_elems2, cut_elems3, cut_elems4]] @@ -490,7 +516,7 @@ def prepare_compute_cuts(self, graph_mock, fw_impl_mock, fw_info_mock, mocker): return ru_calc, cuts, nodes def test_get_cut_target_nodes(self, prepare_compute_cuts): - ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2) = prepare_compute_cuts + ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2, qp) = prepare_compute_cuts assert len(TIC) == 4 sorted_res = lambda res: sorted(res, key=lambda n: n.name) assert sorted_res(ru_calc._get_cut_target_nodes(cut2, TIC.Any)) == [mp, mp_reuse, noq, sp] @@ -499,7 +525,7 @@ def test_get_cut_target_nodes(self, prepare_compute_cuts): assert sorted_res(ru_calc._get_cut_target_nodes(cut2, TIC.QNonConfigurable)) == [sp] def test_compute_act_utilization_by_cut(self, prepare_compute_cuts): - ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2) = prepare_compute_cuts + ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2, qp) = prepare_compute_cuts ru_calc.compute_node_activation_tensor_utilization = Mock(wraps=ru_calc.compute_node_activation_tensor_utilization) ru_calc._get_cut_target_nodes = Mock(wraps=ru_calc._get_cut_target_nodes) @@ -509,13 +535,13 @@ def test_compute_act_utilization_by_cut(self, prepare_compute_cuts): total, per_cut, per_cut_node = ru_calc.compute_activation_utilization_by_cut(TIC.AnyQuantized, BM.QCustom, qcs) cut_nodes_calls = ru_calc._get_cut_target_nodes.call_args_list - assert len(cut_nodes_calls ) == 4 + assert len(cut_nodes_calls) == 4 assert {call.args[0] for call in cut_nodes_calls} == {cut1, cut2, cut3, cut4} assert {call.args[1] for call in cut_nodes_calls } == {TIC.AnyQuantized} compute_tensor_calls = sorted(ru_calc.compute_node_activation_tensor_utilization.call_args_list, key=lambda call: call.args[0].name) - assert len(compute_tensor_calls) == 6 + assert len(compute_tensor_calls) == 7 assert compute_tensor_calls[0].args == (mp, TIC.AnyQuantized, BM.QCustom, qcs['mp']) assert compute_tensor_calls[-1].args == (sp, TIC.AnyQuantized, BM.QCustom, None) @@ -525,17 +551,17 @@ def test_compute_act_utilization_by_cut(self, prepare_compute_cuts): 'mp': Utilization(50, 62.5), 'sp': Utilization(200, 75.)} assert per_cut_node[cut3] == {'sp': Utilization(200, 75.)} - assert per_cut_node[cut4] == {'mp2': Utilization(150, 600.)} + assert per_cut_node[cut4] == {'mp2': Utilization(150, 600.), 'qp': Utilization(150, 600.)} assert per_cut == {cut1: Utilization(24, 21.), cut2: Utilization(274, 158.5), cut3: Utilization(200, 75.), - cut4: Utilization(150, 600.), + cut4: Utilization(300, 1200.), } - assert total == 600. + assert total == 1200. def test_compute_act_utilization_by_cut_no_cut_nodes(self, prepare_compute_cuts): - ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2) = prepare_compute_cuts + ru_calc, (cut1, cut2, cut3, cut4), (mp_reuse, mp, noq, sp, mp2, qp) = prepare_compute_cuts total, per_cut, per_cut_node = ru_calc.compute_activation_utilization_by_cut(TIC.QNonConfigurable, BM.QDefaultSP) assert len(per_cut_node) == 2 @@ -590,13 +616,11 @@ def test_get_w_nbits(self, graph_mock, fw_impl_mock, fw_info_mock): assert ru_calc._get_weight_nbits(node, 'mp', BM.Float, w_qc=None) == FLOAT_BITWIDTH assert ru_calc._get_weight_nbits(node, 'mp', BM.QMinBit, w_qc=None) == 2 assert ru_calc._get_weight_nbits(node, 'mp', BM.QMaxBit, w_qc=None) == 7 - assert ru_calc._get_weight_nbits(node, 'mp', BM.Q8Bit, w_qc=None) == 8 # non-configurable attr with multiple qcs with same w precision - for bm in set(BitwidthMode) - {BM.Float, BM.Q8Bit}: + for bm in set(BitwidthMode) - {BM.Float}: assert ru_calc._get_weight_nbits(node, 'sp', bm, w_qc=None) == 5 assert ru_calc._get_weight_nbits(node, 'sp', BM.Float, w_qc=None) == FLOAT_BITWIDTH - assert ru_calc._get_weight_nbits(node, 'sp', BM.Q8Bit, w_qc=None) == 8 # positional assert ru_calc._get_weight_nbits(node, 2, BM.QMaxBit, w_qc=None) == 6 @@ -924,7 +948,6 @@ def test_compute_regular_node_bops(self, fw_impl_mock, fw_info_mock): ru_calc = ResourceUtilizationCalculator(graph, fw_impl_mock, fw_info_mock) assert ru_calc.compute_node_bops(n2, TIC.AnyQuantized, BM.Float) == 42*32*32 - assert ru_calc.compute_node_bops(n2, TIC.AnyQuantized, BM.Q8Bit) == 42*8*8 assert ru_calc.compute_node_bops(n2, TIC.AnyQuantized, BM.QMinBit) == 42*7*2 assert ru_calc.compute_node_bops(n2, TIC.AnyQuantized, BM.QMaxBit) == 42*16*4 @@ -990,12 +1013,10 @@ def test_compute_virtual_aw_node_bops_fully_quantized(self, fw_impl_mock, fw_inf ru_calc = ResourceUtilizationCalculator(g, fw_impl_mock, fw_info_mock) assert ru_calc.compute_node_bops(a1w2, TIC.AnyQuantized, BM.Float) == 42 * 32 * 32 - assert ru_calc.compute_node_bops(a1w2, TIC.AnyQuantized, BM.Q8Bit) == 42 * 8 * 8 assert ru_calc.compute_node_bops(a1w2, TIC.AnyQuantized, BM.QMinBit) == 42 * 4 * 6 assert ru_calc.compute_node_bops(a1w2, TIC.AnyQuantized, BM.QMaxBit) == 42 * 16 * 16 assert ru_calc.compute_node_bops(a2w3, TIC.AnyQuantized, BM.Float) == 142 * 32 * 32 - assert ru_calc.compute_node_bops(a2w3, TIC.AnyQuantized, BM.Q8Bit) == 142 * 8 * 8 assert ru_calc.compute_node_bops(a2w3, TIC.AnyQuantized, BM.QMinBit) == 142 * 2 * 2 assert ru_calc.compute_node_bops(a2w3, TIC.AnyQuantized, BM.QMaxBit) == 142 * 5 * 6 diff --git a/tests_pytest/conftest.py b/tests_pytest/conftest.py index 3260277b6..466d99498 100644 --- a/tests_pytest/conftest.py +++ b/tests_pytest/conftest.py @@ -31,7 +31,7 @@ def minimal_tpc(): @fixture def graph_mock(): """ Basic Graph mock. """ - return Mock(spec_set=Graph, nodes=[], retrieve_preserved_quantization_node=lambda x: x) + return Mock(spec_set=Graph, nodes=[]) @fixture