def test_forward_trace_functor(): from nncf.dynamic_graph.patch_pytorch import ForwardTraceOnly from nncf.dynamic_graph.trace_tensor import TracedTensor, TensorMeta func = ForwardTraceOnly() shape1, shape2 = ([32, 1, 4, 8], [1, 8, 12, 16]) meta1, meta2 = (TensorMeta(5, 1, shape1), TensorMeta(3, 8, shape2)) input_tensor1 = TracedTensor.from_torch_tensor(torch.Tensor(shape1), meta1) input_tensor2 = TracedTensor.from_torch_tensor(torch.Tensor(shape2), meta2) # 1 -> 1 output_tensor = func(torch.Tensor.view, input_tensor1, [-1]) assert output_tensor.tensor_meta == input_tensor1.tensor_meta # 1 -> N outputs = func(torch.Tensor.chunk, input_tensor1, 3) for out in outputs: assert out.tensor_meta == input_tensor1.tensor_meta # N -> N (2 -> 2) outputs = func(lambda x: x + [5], [input_tensor1, input_tensor2]) assert outputs[0].tensor_meta == input_tensor1.tensor_meta assert outputs[1].tensor_meta == input_tensor2.tensor_meta # M -> N (2 -> 3) with pytest.raises(RuntimeError): outputs = func(lambda x: x + [torch.Tensor(shape2)], [input_tensor1, input_tensor2]) # M -> N (2 -> 1) with pytest.raises(RuntimeError): outputs = func(lambda x: x[0], [input_tensor1, input_tensor2])
def make_op_exec_context_for_coalescing_test(scope_str: str) -> OperationExecutionContext: ia_op_exec_context = InputAgnosticOperationExecutionContext.from_str(scope_str) op_exec_context = OperationExecutionContext(ia_op_exec_context.operator_name, ia_op_exec_context.scope_in_model, ia_op_exec_context.call_order, [TensorMeta(0, 0, [1])]) return op_exec_context
def __call__(self, lhs: TensorMeta, rhs: TensorMeta) -> bool: return TensorMeta.default_comparator(lhs, rhs)